#!/usr/bin/env python3
"""Ordinal-local signing and SCM mediation boundary for a model worker."""

from __future__ import annotations

import json
import os
import threading
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler
from pathlib import Path
from collections.abc import Callable
from typing import Any

import cli_lane_goal
from execution_pool_protocol import (
    MAX_WIRE_BYTES,
    PROTOCOL_VERSION,
    BoundedHTTPServer,
    ProtocolError,
    canonical_json,
    parse_wire,
    read_key,
    sign_envelope,
    verify_envelope,
)
from execution_pool_scm import Boundary as SCMBoundary


KEY_PATH = Path(
    os.environ.get("HERMES_EXECUTION_POOL_KEY_FILE", "/pool-access/execution-pool-key")
)
COORDINATOR = os.environ.get(
    "HERMES_EXECUTION_POOL_URL",
    "http://hermes-execution-pool.hermes.svc.cluster.local:9007",
).rstrip("/")
ORDINAL = int(os.environ.get("HERMES_WORKER_ORDINAL", "-1"))
PORT = int(os.environ.get("HERMES_EXECUTION_CLIENT_PORT", "9009"))
RESULT_FIELDS = frozenset(
    {"status", "summary", "changed_files", "tests_run", "artifacts", "findings", "blockers"}
)


def _binding(value: dict[str, Any]) -> dict[str, Any]:
    return {
        name: value[name]
        for name in ("board", "task_id", "run_id", "worker_ordinal", "attempt")
    }


def _validate_result(payload: Any) -> dict[str, Any]:
    if not isinstance(payload, dict):
        raise ProtocolError("terminal result payload must be an object")
    structured = payload.get("structured")
    if not isinstance(structured, dict) or set(structured) != RESULT_FIELDS:
        raise ProtocolError("terminal result fields do not match the reviewed schema")
    if structured.get("status") not in cli_lane_goal.RESULT_STATUSES:
        raise ProtocolError("terminal result status is invalid")
    if not isinstance(structured.get("summary"), str) or not structured["summary"].strip():
        raise ProtocolError("terminal result summary is required")
    for name in RESULT_FIELDS - {"status", "summary"}:
        value = structured.get(name)
        if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
            raise ProtocolError(f"terminal result {name} must be a text list")
    return payload


class ClientBoundary:
    """Keep HMAC and SCM authority outside the model-facing container."""

    def __init__(self, key: bytes, scm: SCMBoundary | None = None):
        self.key = key
        self.scm = scm or SCMBoundary(key)
        self.current: dict[str, Any] | None = None
        self.resumed: dict[str, str] | None = None
        self.resumed_evidence: dict[str, Any] | None = None
        self.resume_transient = False
        self.lock = threading.RLock()

    def _post(self, path: str, envelope: dict[str, Any]) -> dict[str, Any]:
        request = urllib.request.Request(
            COORDINATOR + path,
            data=canonical_json(envelope),
            method="POST",
            headers={"Content-Type": "application/json", "Cache-Control": "no-store"},
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            body = response.read(MAX_WIRE_BYTES + 1)
        if len(body) > MAX_WIRE_BYTES:
            raise ProtocolError("coordinator response exceeds the wire limit")
        return verify_envelope(self.key, json.loads(body))

    def poll(self) -> dict[str, Any]:
        poll_binding = {
            "board": "",
            "task_id": "",
            "run_id": "",
            "worker_ordinal": ORDINAL,
            "attempt": 0,
        }
        response = self._post(
            "/v1/poll",
            sign_envelope(self.key, "poll", poll_binding, {"ready": True}),
        )
        with self.lock:
            if response["kind"] == "ack":
                self.current = None
                return {"assignment": None}
            if response["kind"] != "assignment" or response["worker_ordinal"] != ORDINAL:
                raise ProtocolError("coordinator returned a foreign assignment")
            publication_only = isinstance(response.get("payload", {}).get("scm_resume"), dict)
            checkout = {"workspace": "", "baseline_sha": ""} if publication_only else self.scm.checkout(response)
            self.current = response
            self.resumed = None
            self.resumed_evidence = None
            self.resume_transient = False
            assignment = {
                **_binding(response),
                "payload": response["payload"],
                "workspace": checkout["workspace"],
                "baseline_sha": checkout["baseline_sha"],
                "protocol_version": PROTOCOL_VERSION,
            }
        return {"assignment": assignment}

    def _current_for(self, supplied: Any) -> tuple[dict[str, Any], dict[str, Any]]:
        if not isinstance(supplied, dict):
            raise ProtocolError("local request binding must be an object")
        if self.current is None or supplied != _binding(self.current):
            raise ProtocolError("local request does not own the current assignment")
        return self.current, dict(supplied)

    def _publication_lease(self, binding: dict[str, Any]) -> tuple[Callable[[], None], Callable[[], None]]:
        """Keep a resume lease alive outside the SCM lock and fail before another step."""
        stop, failures = threading.Event(), []

        def renew() -> None:
            try:
                response = self._post("/v1/heartbeat", sign_envelope(self.key, "heartbeat", binding, {"note": "SCM publication in progress"}))
                if response["kind"] != "ack" or _binding(response) != binding or not response["payload"].get("accepted"):
                    raise ProtocolError("coordinator rejected SCM publication lease")
            except (OSError, ProtocolError, ValueError, urllib.error.URLError) as error:
                failures.append(error)

        renew()
        def keepalive() -> None:
            while not stop.wait(20):
                renew()

        thread = threading.Thread(target=keepalive, daemon=True)
        thread.start()

        def checkpoint() -> None:
            if failures:
                raise ProtocolError("SCM publication lease was lost") from failures[0]

        return checkpoint, lambda: (stop.set(), thread.join(timeout=1))

    def heartbeat(self, request: dict[str, Any]) -> dict[str, Any]:
        payload = request.get("payload")
        if not isinstance(payload, dict):
            raise ProtocolError("heartbeat payload must be an object")
        with self.lock:
            _assignment, binding = self._current_for(request.get("binding"))
            response = self._post(
                "/v1/heartbeat",
                sign_envelope(self.key, "heartbeat", binding, payload),
            )
            if response["kind"] != "ack" or _binding(response) != binding:
                raise ProtocolError("coordinator acknowledgement binding changed")
        return {"ack": response["payload"]}

    def _submit(
        self, assignment: dict[str, Any], request: dict[str, Any],
        structured: dict[str, Any],
    ) -> dict[str, str] | None:
        """Publish the run's work, downgrading the result rather than losing it.

        A refused push or draft used to unwind the whole worker, so the run's
        actual output was discarded and the task was blocked with only the SCM
        error. The commits live on the durable workspace volume, so the result is
        still reported here -- as blocked, carrying the reason -- which keeps the
        prior work recoverable and reviewable by a human.
        """
        try:
            if isinstance(self.scm, SCMBoundary):
                checkpoint, close = self._publication_lease(_binding(assignment))
                try:
                    submission = self.scm.submit(assignment, request, checkpoint=checkpoint)
                    checkpoint()
                finally:
                    close()
            else:
                submission = self.scm.submit(assignment, request)
        except (OSError, ProtocolError, RuntimeError, ValueError) as error:
            payload = assignment.get("payload")
            resume = None
            if isinstance(payload, dict) and payload.get("continuation_kind") == "repair":
                try:
                    resume = self.scm.resume_artifact(assignment, request, structured)
                except (OSError, ProtocolError, RuntimeError, ValueError):
                    resume = None
            structured["status"] = "blocked"
            reason = (
                "Distributed SCM submission failed; the commits remain on this "
                f"ordinal's workspace: {type(error).__name__}: {error}"
            )
            if reason not in structured["blockers"]:
                structured["blockers"].append(reason)
            return {"resume": resume} if resume is not None else None
        pull = str(submission.get("pull_request") or "")
        branch = str(submission.get("branch") or "")
        for artifact in (pull, f"branch:{branch}" if branch else ""):
            if artifact and artifact not in structured["artifacts"]:
                structured["artifacts"].append(artifact)
        head = str(submission.get("head") or "")
        return {"branch": branch, "pull_request": pull, "head": head}

    def resume(self, request: dict[str, Any]) -> dict[str, Any]:
        """Publish the coordinator-carried artifact without exposing it to the worker."""
        with self.lock:
            assignment, _binding = self._current_for(request.get("binding"))
            artifact = assignment.get("payload", {}).get("scm_resume")
            if not isinstance(artifact, dict):
                raise ProtocolError("assignment has no SCM resume artifact")
            evidence = _validate_result({"structured": artifact.get("structured")})["structured"]
            if evidence["status"] != "completed":
                raise ProtocolError("SCM resume evidence is not completed")
            checkpoint, close = self._publication_lease(_binding)
            try:
                self.resume_transient = False
                try:
                    self.resumed = self.scm.resume(assignment, artifact, checkpoint=checkpoint)
                    checkpoint()
                except (OSError, RuntimeError, urllib.error.URLError) as error:
                    detail = str(error).lower()
                    if any(marker in detail for marker in ("timed out", "connection", "could not resolve", "http 502", "http 503", "http 504")):
                        self.resume_transient = True
                        return {"publication_retry_transient": True}
                    raise ProtocolError("SCM resume publication was rejected") from error
            finally:
                close()
            self.resumed_evidence = json.loads(canonical_json(evidence))
            return {"scm_submission": self.resumed}

    def finish(self, request: dict[str, Any]) -> dict[str, Any]:
        payload = _validate_result(request.get("payload"))
        # The model-facing caller cannot classify a retry as transient.  Only a
        # preceding mediator resume may attach this coordinator control signal.
        payload.pop("publication_retry_transient", None)
        with self.lock:
            assignment, binding = self._current_for(request.get("binding"))
            structured = payload["structured"]
            if self.resume_transient:
                payload["publication_retry_transient"] = True
            if structured["status"] == "completed" and int(payload.get("returncode", 1)) == 0:
                if self.resumed is not None:
                    if self.resumed_evidence is None:
                        raise ProtocolError("SCM resume evidence is unavailable")
                    structured = json.loads(canonical_json(self.resumed_evidence))
                    payload["structured"] = structured
                    payload["returncode"] = 0
                    payload["scm_submission"] = self.resumed
                else:
                    submission = self._submit(assignment, request, structured)
                    if submission is not None and "resume" not in submission:
                        # The mediator, not the model-facing request, records the
                        # broker-confirmed PR/branch in the signed terminal wire.
                        payload["scm_submission"] = submission
                    elif submission is not None:
                        # The exact clean commit remains on the durable worker volume.
                        # A broker refusal is execution infrastructure, not a provider
                        # capability verdict, so let the coordinator retry it after the
                        # scoped SCM condition is corrected.
                        payload["capacity_failure"] = True
                        payload["scm_resume"] = submission["resume"]
                    else:
                        payload["capacity_failure"] = True
            response = self._post(
                "/v1/result", sign_envelope(self.key, "result", binding, payload)
            )
            if response["kind"] != "ack" or _binding(response) != binding:
                raise ProtocolError("coordinator acknowledgement binding changed")
            if response["payload"].get("accepted"):
                self.current = None
        return {"ack": response["payload"], "structured": structured}


def handler_factory(boundary: ClientBoundary) -> type[BaseHTTPRequestHandler]:
    class Handler(BaseHTTPRequestHandler):
        server_version = f"hermes-execution-mediator/{PROTOCOL_VERSION}"

        def _reply(self, status: int, value: dict[str, Any]) -> None:
            body = canonical_json(value)
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self) -> None:  # noqa: N802
            value = {"ready": True, "protocol_version": PROTOCOL_VERSION}
            self._reply(200, value) if self.path == "/ready" else self._reply(
                404, {"error": "not found"}
            )

        def do_POST(self) -> None:  # noqa: N802
            try:
                length = int(self.headers.get("Content-Length", "0"))
                request = (
                    parse_wire(self.rfile.read(length))
                    if 0 < length <= MAX_WIRE_BYTES
                    else None
                )
                allowed = {"operation", "binding", "payload", "title", "body"}
                if not request or set(request) - allowed:
                    raise ProtocolError("invalid local mediator request")
                operation = str(request.get("operation") or "")
                routes = {
                    "poll": boundary.poll,
                    "heartbeat": lambda: boundary.heartbeat(request),
                    "resume": lambda: boundary.resume(request),
                    "finish": lambda: boundary.finish(request),
                }
                if operation not in routes:
                    raise ProtocolError("unsupported local mediator operation")
                self._reply(200, routes[operation]())
            except (ProtocolError, OSError, ValueError, urllib.error.URLError) as error:
                self._reply(409, {"error": str(error)[:2000]})

        def log_message(self, _format: str, *_arguments: Any) -> None:
            return

    return Handler


def main() -> int:
    if ORDINAL not in range(3):
        raise SystemExit("HERMES_WORKER_ORDINAL must be 0, 1, or 2")
    key = read_key(KEY_PATH)
    BoundedHTTPServer(
        ("0.0.0.0", PORT),
        handler_factory(ClientBoundary(key)),
        max_workers=4,
    ).serve_forever()
    return 0


if __name__ == "__main__":  # pragma: no cover - exercised by the container entrypoint
    raise SystemExit(main())