"""Project-local HTTP/1.1 transport for Responses-compatible model gateways."""
from contextlib import contextmanager
from collections import deque
import http.server
import http.client
import json
from pathlib import Path
import re
import socket
import threading
import time
import urllib.error
import urllib.request
from urllib.parse import urlsplit


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *_args, **_kwargs):
        return None


class RequestCancelled(OSError):
    pass


class ModelTransport:
    def __init__(self, provider, upstream, *, timeout=60, retries=1, retry_delay=0.5, log_path=None,
                 max_request_bytes=512 * 1024):
        parsed = urlsplit(upstream)
        local_http = parsed.scheme == "http" and parsed.hostname in {"127.0.0.1", "localhost"}
        if (not isinstance(provider, str) or not re.fullmatch(r"[A-Za-z0-9_-]+", provider)
                or not parsed.hostname or parsed.scheme != "https" and not local_http
                or parsed.username or parsed.password or parsed.query or parsed.fragment):
            raise ValueError("model transport requires a named provider and a credential-free HTTPS upstream")
        self.provider = provider
        self.upstream = upstream.rstrip("/")
        self.timeout = float(timeout)
        if not 1 <= self.timeout <= 300:
            raise ValueError("model transport timeout must be between 1 and 300 seconds")
        self.retries = int(retries)
        if not 0 <= self.retries <= 2:
            raise ValueError("model transport retries must be between 0 and 2")
        self.retry_delay = float(retry_delay)
        if not 0 <= self.retry_delay <= 5:
            raise ValueError("model transport retry delay must be between 0 and 5 seconds")
        self.log_path = Path(log_path) if log_path else None
        self.max_request_bytes = int(max_request_bytes)
        if not 1024 <= self.max_request_bytes <= 64 * 1024 * 1024:
            raise ValueError("model request byte budget must be between 1024 and 67108864")
        self.started = time.monotonic()
        self.lock = threading.Lock()
        self.slots = threading.BoundedSemaphore(4)
        self.requests = 0
        self.request_usage = deque(maxlen=256)
        self.active = {}
        self.finished = deque(maxlen=256)
        self.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), self.handler())
        self.server.daemon_threads = True
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()

    def event(self, event, **fields):
        if self.log_path is None:
            return
        record = {"event": event, "seconds": round(time.monotonic() - self.started, 3), **fields}
        with self.lock:
            self.log_path.parent.mkdir(parents=True, exist_ok=True)
            with self.log_path.open("a", encoding="utf-8") as stream:
                stream.write(json.dumps(record) + "\n")

    @property
    def url(self):
        return f"http://127.0.0.1:{self.server.server_port}"

    def usage_checkpoint(self):
        with self.lock:
            return self.requests

    def begin_request(self):
        with self.lock:
            self.requests += 1
            number = self.requests
            now = time.monotonic()
            self.active[number] = {"request": number, "started": now, "lastData": now,
                                   "phase": "request", "bytes": 0, "cancel": threading.Event()}
            return number

    def update_request(self, number, **values):
        with self.lock:
            if number in self.active:
                self.active[number].update(values)

    def ensure_active(self, number):
        with self.lock:
            state = self.active.get(number)
            if state and state["cancel"].is_set():
                raise RequestCancelled("model request was cancelled")

    def finish_request(self, number):
        with self.lock:
            state = self.active.pop(number, None)
            if state:
                if state["cancel"].is_set():
                    state["phase"] = "cancelled"
                self.finished.append({key: value for key, value in state.items()
                                      if key not in {"cancel", "connection", "response"}})

    def activity_since(self, checkpoint):
        now = time.monotonic()
        with self.lock:
            rows = [row for row in [*self.finished, *self.active.values()] if row["request"] > checkpoint]
            if not rows:
                return {}
            latest = max(rows, key=lambda row: row["request"])
            elapsed, idle = now - latest["started"], now - latest["lastData"]
            phase = latest["phase"]
            # Heartbeats cannot extend the absolute request budget. A small
            # completed-response grace lets the CLI dispatch buffered events.
            if phase == "request":
                waiting = elapsed < self.response_budget
            elif phase in {"headers", "streaming"}:
                waiting = elapsed < self.response_budget and idle < self.timeout
            else:
                waiting = phase == "completed" and idle < 5
            return {"request": latest["request"], "phase": phase, "bytes": latest["bytes"],
                    "requestBytes": latest.get("requestBytes"), "errorType": latest.get("errorType"),
                    "elapsedSeconds": round(elapsed, 3), "idleSeconds": round(idle, 3),
                    "responseBudgetSeconds": self.response_budget, "canWait": bool(waiting),
                    "activeRequests": sum(row["request"] > checkpoint for row in self.active.values())}

    def cancel_since(self, checkpoint):
        with self.lock:
            rows = [row.copy() for row in self.active.values() if row["request"] > checkpoint]
            for row in rows:
                row["cancel"].set()
                self.active[row["request"]]["phase"] = "cancelled"
        for row in rows:
            connection, response = row.get("connection"), row.get("response")
            sockets = [getattr(connection, "sock", None), getattr(getattr(getattr(response, "fp", None), "raw", None), "_sock", None)]
            for stream in sockets:
                if stream is not None:
                    try:
                        stream.shutdown(socket.SHUT_RDWR)
                    except OSError:
                        pass
            if connection is not None:
                connection.close()
            self.event("cancelled", request=row["request"])
        return len(rows)

    def record_usage(self, number, usage):
        if not isinstance(usage, dict):
            return
        tokens = usage.get("input_tokens")
        if type(tokens) is not int or tokens < 0:
            return
        with self.lock:
            self.request_usage.append((number, tokens))
        self.event("usage", request=number, inputTokens=tokens)

    def usage_since(self, checkpoint):
        with self.lock:
            rows = sorted((number, tokens) for number, tokens in self.request_usage if number > checkpoint)
            requests = self.requests - checkpoint
        return {"requests": requests, "measuredRequests": len(rows),
                "lastInputTokens": rows[-1][1] if rows else None,
                "peakInputTokens": max((tokens for _, tokens in rows), default=None),
                "complete": requests > 0 and len(rows) == requests}

    @property
    def response_budget(self):
        return (self.retries + 1) * self.timeout + self.retry_delay * self.retries * (self.retries + 1) / 2 + 10

    def cli_args(self):
        return ["-c", f"model_provider={json.dumps(self.provider)}",
                "-c", f"model_providers.{self.provider}.base_url={json.dumps(self.url)}",
                "-c", "features.enable_request_compression=false"]

    def handler(self):
        transport = self

        class Handler(http.server.BaseHTTPRequestHandler):
            protocol_version = "HTTP/1.1"

            def log_message(self, *_):
                pass

            def respond(self, status, body, content_type="application/json", retry_after=None):
                self.send_response(status)
                self.send_header("Content-Type", content_type)
                self.send_header("Content-Length", str(len(body)))
                self.send_header("Connection", "close")
                if retry_after:
                    self.send_header("Retry-After", retry_after)
                self.end_headers()
                self.wfile.write(body)
                self.wfile.flush()

            def do_POST(self):
                if self.path not in {"/responses", "/responses/compact"}:
                    self.send_error(404)
                    return
                if not transport.slots.acquire(blocking=False):
                    self.respond(503, b'{"error":{"message":"Local model transport is busy"}}')
                    return
                try:
                    self.forward()
                finally:
                    self.close_connection = True
                    transport.slots.release()

            def forward(self):
                sent_headers = False
                number = transport.begin_request()
                try:
                    self.connection.settimeout(30)
                    length = int(self.headers.get("Content-Length", "0"))
                    transport.update_request(number, requestBytes=length)
                    if not 0 < length <= 64 * 1024 * 1024 or self.headers.get("Transfer-Encoding"):
                        self.respond(413, b'{"error":{"message":"Invalid model request size or framing"}}')
                        return
                    if length > transport.max_request_bytes:
                        transport.update_request(number, phase="context_rejected")
                        transport.event("context_rejected", request=number, bytes=length, limit=transport.max_request_bytes)
                        self.respond(413, b'{"error":{"message":"Model request exceeds local context byte budget; checkpoint and start a fresh bounded turn","code":"context_budget_exceeded"}}')
                        return
                    body = self.rfile.read(length)
                    if len(body) != length:
                        self.respond(400, b'{"error":{"message":"Incomplete model request"}}')
                        return
                    excluded = {"host", "connection", "content-length", "transfer-encoding", "keep-alive", "te", "trailer",
                                "upgrade", "proxy-authorization", "proxy-connection", "accept-encoding"}
                    headers = {k: v for k, v in self.headers.items() if k.lower() not in excluded}
                    headers.update({"Accept-Encoding": "identity", "Connection": "close"})
                    transport.event("request", request=number, path=self.path, bytes=length)
                    if not self.headers.get("Content-Encoding"):
                        try:
                            payload = json.loads(body)
                            inputs = payload.get("input", [])
                            tools = payload.get("tools", [])
                            tool_outputs = [row.get("output", "") for row in inputs
                                            if isinstance(row, dict) and row.get("type") in {"function_call_output", "custom_tool_call_output"}] if isinstance(inputs, list) else []
                            transport.event("context", request=number, inputItems=len(inputs) if isinstance(inputs, list) else 1,
                                            tools=len(tools) if isinstance(tools, list) else 0,
                                            toolOutputBytes=sum(len((value if isinstance(value, str) else json.dumps(value)).encode("utf-8")) for value in tool_outputs))
                        except (ValueError, AttributeError):
                            pass
                    request = urllib.request.Request(transport.upstream + self.path, data=body, headers=headers)
                    with transport.open_request(request, number) as response:
                        transport.ensure_active(number)
                        transport.update_request(number, phase="headers", response=response, lastData=time.monotonic())
                        transport.event("headers", request=number, status=response.status)
                        self.send_response(response.status)
                        self.send_header("Content-Type", response.headers.get("Content-Type", "text/event-stream"))
                        self.send_header("Connection", "close")
                        self.end_headers()
                        sent_headers = True
                        size, terminal = 0, None
                        for line in response:
                            if size == 0:
                                transport.event("first_body", request=number)
                            size += len(line)
                            if line.startswith(b"data:"):
                                try:
                                    item = json.loads(line[5:])
                                    transport.update_request(number, phase="streaming", lastData=time.monotonic())
                                    if item.get("type") in {"response.completed", "response.failed", "error"}:
                                        terminal = item["type"]
                                    if item.get("type") == "response.completed" and self.path == "/responses":
                                        transport.record_usage(number, (item.get("response") or {}).get("usage"))
                                except (ValueError, AttributeError):
                                    pass
                            self.wfile.write(line)
                            self.wfile.flush()
                            transport.ensure_active(number)
                            transport.update_request(number, bytes=size)
                        transport.event("end", request=number, bytes=size, terminal=terminal)
                        transport.update_request(number, phase="completed" if terminal == "response.completed" else "incomplete",
                                                 lastData=time.monotonic(), terminal=terminal)
                except urllib.error.HTTPError as error:
                    transport.update_request(number, phase="http_error")
                    transport.event("http_error", request=number, status=error.code)
                    try:
                        self.respond(error.code, error.read(1024 * 1024), error.headers.get("Content-Type", "application/json"),
                                     error.headers.get("Retry-After"))
                    except OSError:
                        pass
                except (OSError, ValueError, http.client.HTTPException) as error:
                    transport.update_request(number, phase="cancelled" if isinstance(error, RequestCancelled) else "error", errorType=type(error).__name__)
                    transport.event("error", request=number, errorType=type(error).__name__, afterHeaders=sent_headers)
                    if not sent_headers:
                        try:
                            self.respond(502, b'{"error":{"message":"Model upstream request failed","type":"upstream_error"}}')
                        except OSError:
                            pass
                finally:
                    transport.finish_request(number)

        return Handler

    def open_request(self, request, number):
        transport = self
        def connection_factory(kind):
            def create(*args, **kwargs):
                transport.ensure_active(number)
                connection = kind(*args, **kwargs)
                transport.update_request(number, connection=connection)
                return connection
            return create
        class HttpHandler(urllib.request.HTTPHandler):
            def http_open(self, request):
                return self.do_open(connection_factory(http.client.HTTPConnection), request)
        class HttpsHandler(urllib.request.HTTPSHandler):
            def https_open(self, request):
                return self.do_open(connection_factory(http.client.HTTPSConnection), request, context=self._context)
        opener = urllib.request.build_opener(NoRedirect(), HttpHandler(), HttpsHandler())
        for attempt in range(self.retries + 1):
            self.ensure_active(number)
            self.event("attempt", request=number, attempt=attempt + 1)
            try:
                return opener.open(request, timeout=self.timeout)
            except urllib.error.HTTPError as error:
                self.ensure_active(number)
                if error.code not in {408, 502, 503, 504} or error.headers.get("Retry-After") or attempt == self.retries:
                    raise
                error.close()
                reason = f"http_{error.code}"
            except (urllib.error.URLError, OSError, http.client.HTTPException) as error:
                self.ensure_active(number)
                if attempt == self.retries:
                    raise
                reason = type(error).__name__
            # No response bytes have reached the client, so no tool call
            # from this attempt can have been dispatched locally.
            self.event("retry", request=number, attempt=attempt + 1, reason=reason)
            with self.lock:
                state = self.active.get(number)
            if state:
                if state["cancel"].wait(self.retry_delay * (attempt + 1)):
                    raise RequestCancelled("cancelled before retry")
            else:
                time.sleep(self.retry_delay * (attempt + 1))

    def close(self):
        self.cancel_since(0)
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=2)


@contextmanager
def configured_transport(root, *, process_id):
    path = Path(root) / "docs/current/scheduler-transport.json"
    if not path.is_file():
        yield None
        return
    config = json.loads(path.read_text(encoding="utf-8"))
    if not config.get("enabled", False):
        yield None
        return
    transport = ModelTransport(config["provider"], config["upstreamBaseUrl"], timeout=config.get("requestTimeoutSeconds", 60),
                               retries=config.get("retriesBeforeHeaders", 1),
                               max_request_bytes=config.get("maxRequestBytes", 512 * 1024),
                               log_path=Path(root) / f"docs/current/continuation-runs/transport-{process_id}.jsonl")
    try:
        yield transport
    finally:
        transport.close()