#!/usr/bin/env python3
"""Route serial-query requests using a fixed local cost calibration."""

import argparse
import json
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from urllib.error import URLError
from urllib.request import Request, urlopen
from urllib.parse import parse_qs, urlsplit


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("calibration", type=Path)
    parser.add_argument("--policy", choices=("user-first", "network", "profiled"),
                        default="profiled")
    parser.add_argument("--compute-a-url", required=True)
    parser.add_argument("--compute-b-url", required=True)
    parser.add_argument("--port", type=int, default=18083)
    args = parser.parse_args()
    calibration = json.loads(args.calibration.read_text())
    targets = {"A": args.compute_a_url, "B": args.compute_b_url}

    class Proxy(BaseHTTPRequestHandler):
        def do_GET(self):
            try:
                count = int(parse_qs(urlsplit(self.path).query).get("count", [""])[0])
                if not 0 <= count <= 64:
                    raise ValueError("count outside 0..64")
            except ValueError as error:
                self.send_error(400, str(error))
                return

            def cost(site):
                row = calibration[site]
                if args.policy == "user-first":
                    return row["health_p50_ms"]
                baseline_ms = (row["baseline_p50_ms"] if args.policy == "profiled"
                               else row["health_p50_ms"])
                return baseline_ms + count * row["db_rtt_p50_ms"]

            site = min(("A", "B"), key=cost)
            try:
                with urlopen(Request(targets[site] + self.path), timeout=10) as response:
                    body, status = response.read(), response.status
            except (OSError, URLError) as error:
                self.send_error(502, str(error))
                return
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("X-Policy", args.policy)
            self.send_header("X-Selected-Host", site)
            self.send_header("X-Estimated-Cost-Ms", f"{cost(site):.3f}")
            self.end_headers()
            self.wfile.write(body)

        def log_message(self, format, *values):
            return

    HTTPServer(("127.0.0.1", args.port), Proxy).serve_forever()


if __name__ == "__main__":
    main()