#!/usr/bin/env python3
"""Authenticated, bounded, restart-safe Hermes execution-pool protocol."""

from __future__ import annotations

import hashlib
import hmac
import json
import os
import re
import sqlite3
import stat
import threading
import time
import uuid
from http.server import ThreadingHTTPServer
from pathlib import Path
from typing import Any


MAX_WIRE_BYTES = 64 * 1024
MAX_ACTIVITY_BYTES = 12 * 1024
MAX_CLOCK_SKEW = 30
MAX_ENVELOPE_LIFETIME = 300
PROTOCOL_VERSION = 2
IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
KINDS = frozenset({"poll", "assignment", "heartbeat", "result", "ack"})


class ProtocolError(ValueError):
    """A request failed the authenticated pool contract."""


class BoundedHTTPServer(ThreadingHTTPServer):
    """Bound concurrent requests and slow clients on internal pool channels."""

    daemon_threads = True
    request_queue_size = 8

    def __init__(self, *args: Any, max_workers: int = 8, **kwargs: Any):
        self._slots = threading.BoundedSemaphore(max(1, min(max_workers, 16)))
        super().__init__(*args, **kwargs)

    def get_request(self) -> tuple[Any, Any]:
        request, address = super().get_request()
        request.settimeout(15)
        return request, address

    def process_request(self, request: Any, client_address: Any) -> None:
        if not self._slots.acquire(blocking=False):
            self.shutdown_request(request)
            return
        try:
            super().process_request(request, client_address)
        except Exception:
            self._slots.release()
            raise

    def process_request_thread(self, request: Any, client_address: Any) -> None:
        try:
            super().process_request_thread(request, client_address)
        finally:
            self._slots.release()


def canonical_json(value: Any) -> bytes:
    """Encode one value deterministically for digests and signatures."""
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")


def atomic_json(path: Path, value: dict[str, Any], mode: int = 0o600) -> None:
    """Durably replace one bounded pool document without following symlinks."""
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    if path.parent.is_symlink():
        raise ProtocolError("pool state directory must not be a symlink")
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(temporary, flags, mode)
    try:
        remaining = memoryview(json.dumps(value, indent=2, sort_keys=True).encode() + b"\n")
        while remaining:
            remaining = remaining[os.write(descriptor, remaining) :]
        os.fsync(descriptor)
        os.fchmod(descriptor, mode)
    finally:
        os.close(descriptor)
    try:
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def payload_digest(payload: Any) -> str:
    return hashlib.sha256(canonical_json(payload)).hexdigest()


def derive_ordinal_key(master: bytes, ordinal: int) -> bytes:
    """Derive one cryptographically isolated worker authority from the pool root."""
    if ordinal not in range(3) or not 32 <= len(master) <= 4096:
        raise ProtocolError("pool key derivation input is invalid")
    context = f"hermes-execution-pool-v2:worker:{ordinal}".encode()
    return hmac.new(master, context, hashlib.sha256).hexdigest().encode()


def envelope_ordinal(envelope: Any) -> int:
    """Read only the bounded ordinal needed to select a verification key."""
    if not isinstance(envelope, dict):
        raise ProtocolError("invalid pool message")
    value = envelope.get("worker_ordinal")
    if type(value) is not int or value not in range(3):
        raise ProtocolError("worker binding is outside the pool")
    return value


def read_key(path: Path) -> bytes:
    """Read a private regular file without following a final symlink."""
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
    except OSError as error:
        raise ProtocolError(f"pool key is unavailable: {error}") from error
    try:
        info = os.fstat(descriptor)
        if not stat.S_ISREG(info.st_mode) or info.st_mode & 0o077:
            raise ProtocolError("pool key must be a private regular file")
        value = os.read(descriptor, 4097).strip()
    finally:
        os.close(descriptor)
    if len(value) < 32 or len(value) > 4096:
        raise ProtocolError("pool key length is outside the safe range")
    return value


def _identifier(name: str, value: Any, *, allow_empty: bool = False) -> str:
    if not isinstance(value, str):
        raise ProtocolError(f"invalid {name}")
    text = value
    if allow_empty and not text:
        return text
    if not IDENTIFIER.fullmatch(text):
        raise ProtocolError(f"invalid {name}")
    return text


def sign_envelope(
    key: bytes,
    kind: str,
    binding: dict[str, Any],
    payload: Any,
    *,
    now: int | None = None,
    lifetime: int = 120,
    delivery_id: str | None = None,
) -> dict[str, Any]:
    """Bind a message to the exact Kanban run, ordinal, and attempt."""
    current = int(time.time()) if now is None else int(now)
    lifetime = max(1, min(int(lifetime), MAX_ENVELOPE_LIFETIME))
    envelope = {
        "version": PROTOCOL_VERSION,
        "kind": kind,
        "board": str(binding.get("board") or ""),
        "task_id": str(binding.get("task_id") or ""),
        "run_id": str(binding.get("run_id") or ""),
        "worker_ordinal": int(binding.get("worker_ordinal", -1)),
        "attempt": int(binding.get("attempt", 0)),
        "delivery_id": delivery_id or str(uuid.uuid4()),
        "issued_at": current,
        "expires_at": current + lifetime,
        "payload_digest": payload_digest(payload),
        "payload": payload,
    }
    if kind not in KINDS:
        raise ProtocolError("unsupported message kind")
    unsigned = canonical_json(envelope)
    if len(unsigned) > MAX_WIRE_BYTES:
        raise ProtocolError("pool message exceeds the wire limit")
    envelope["signature"] = hmac.new(key, unsigned, hashlib.sha256).hexdigest()
    return envelope


def verify_envelope(
    key: bytes,
    envelope: Any,
    *,
    expected_kind: str | None = None,
    now: int | None = None,
) -> dict[str, Any]:
    """Verify structure, lifetime, digest, and HMAC before using a message."""
    if not isinstance(envelope, dict) or len(canonical_json(envelope)) > MAX_WIRE_BYTES:
        raise ProtocolError("invalid or oversized pool message")
    required = {"version", "kind", "board", "task_id", "run_id", "worker_ordinal", "attempt", "delivery_id", "issued_at", "expires_at", "payload_digest", "payload", "signature"}
    if set(envelope) != required or envelope.get("version") != PROTOCOL_VERSION:
        raise ProtocolError(
            f"pool message fields do not match version {PROTOCOL_VERSION}"
        )
    kind = envelope["kind"]
    if not isinstance(kind, str):
        raise ProtocolError("unexpected message kind")
    if kind not in KINDS or (expected_kind and kind != expected_kind):
        raise ProtocolError("unexpected message kind")
    empty_binding = kind in {"poll", "ack"}
    _identifier("board", envelope["board"], allow_empty=empty_binding)
    _identifier("task_id", envelope["task_id"], allow_empty=empty_binding)
    _identifier("run_id", envelope["run_id"], allow_empty=empty_binding)
    _identifier("delivery_id", envelope["delivery_id"])
    numeric = tuple(envelope[name] for name in ("worker_ordinal", "attempt", "issued_at", "expires_at"))
    if any(type(value) is not int for value in numeric):
        raise ProtocolError("invalid numeric binding")
    ordinal, attempt, issued, expires = numeric
    if ordinal not in range(3) or attempt < 0:
        raise ProtocolError("worker binding is outside the pool")
    current = int(time.time()) if now is None else int(now)
    if issued > current + MAX_CLOCK_SKEW or expires < current - MAX_CLOCK_SKEW:
        raise ProtocolError("pool message is outside its validity window")
    if expires <= issued or expires - issued > MAX_ENVELOPE_LIFETIME:
        raise ProtocolError("pool message lifetime is invalid")
    if envelope["payload_digest"] != payload_digest(envelope["payload"]):
        raise ProtocolError("pool payload digest mismatch")
    signature = str(envelope["signature"])
    unsigned = dict(envelope)
    unsigned.pop("signature")
    expected = hmac.new(key, canonical_json(unsigned), hashlib.sha256).hexdigest()
    if not hmac.compare_digest(signature, expected):
        raise ProtocolError("pool message authentication failed")
    return envelope


def parse_wire(body: bytes) -> dict[str, Any]:
    if not body or len(body) > MAX_WIRE_BYTES:
        raise ProtocolError("empty or oversized request")
    try:
        value = json.loads(body)
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise ProtocolError("malformed JSON request") from error
    if not isinstance(value, dict):
        raise ProtocolError("request must be a JSON object")
    return value


class PoolStore:
    """Coordinator-owned durable assignments; never stores provider secrets."""

    def __init__(self, path: Path, lease_seconds: int = 90):
        self.path = path
        self.lease_seconds = max(60, min(int(lease_seconds), 600))
        self._lock = threading.RLock()
        path.parent.mkdir(parents=True, exist_ok=True)
        self._initialize()

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA journal_mode=WAL")
        connection.execute("PRAGMA synchronous=FULL")
        connection.execute("PRAGMA busy_timeout=10000")
        return connection

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS assignments (
                  board TEXT NOT NULL, task_id TEXT NOT NULL, run_id TEXT NOT NULL,
                  worker_ordinal INTEGER NOT NULL CHECK(worker_ordinal BETWEEN 0 AND 2),
                  attempt INTEGER NOT NULL, assignment_digest TEXT NOT NULL,
                  payload_json TEXT NOT NULL, state TEXT NOT NULL,
                  lease_until REAL NOT NULL DEFAULT 0, last_heartbeat REAL NOT NULL DEFAULT 0,
                  result_digest TEXT, result_json TEXT, created_at REAL NOT NULL,
                  updated_at REAL NOT NULL, PRIMARY KEY(board, task_id, run_id)
                );
                CREATE UNIQUE INDEX IF NOT EXISTS one_live_assignment_per_worker
                  ON assignments(worker_ordinal) WHERE state IN ('assigned','running','result');
                CREATE TABLE IF NOT EXISTS deliveries (
                  delivery_id TEXT PRIMARY KEY, kind TEXT NOT NULL, digest TEXT NOT NULL,
                  received_at REAL NOT NULL
                );
                """
            )

    @staticmethod
    def _record(row: sqlite3.Row | None) -> dict[str, Any] | None:
        if row is None:
            return None
        value = dict(row)
        value["payload"] = json.loads(value.pop("payload_json"))
        if value.get("result_json"):
            value["result"] = json.loads(value["result_json"])
        return value

    def add(self, binding: dict[str, Any], payload: dict[str, Any]) -> bool:
        """Create exactly one assignment for a claimed run and free ordinal."""
        now = time.time()
        digest = payload_digest(payload)
        values = (
            binding["board"], binding["task_id"], binding["run_id"],
            binding["worker_ordinal"], binding["attempt"], digest,
            canonical_json(payload).decode(), "assigned", now, now,
        )
        with self._lock, self._connect() as connection:
            try:
                connection.execute("BEGIN IMMEDIATE")
                connection.execute(
                    """INSERT INTO assignments
                    (board,task_id,run_id,worker_ordinal,attempt,assignment_digest,
                     payload_json,state,created_at,updated_at)
                    VALUES (?,?,?,?,?,?,?,?,?,?)""",
                    values,
                )
                connection.commit()
                return True
            except sqlite3.IntegrityError as error:
                connection.rollback()
                existing = connection.execute(
                    """SELECT assignment_digest,worker_ordinal,attempt FROM assignments
                    WHERE board=? AND task_id=? AND run_id=?""",
                    values[:3],
                ).fetchone()
                if existing and tuple(existing) == (
                    digest, binding["worker_ordinal"], binding["attempt"]
                ):
                    return False
                if existing:
                    raise ProtocolError("conflicting duplicate assignment") from error
                raise ProtocolError("worker ordinal already has a live assignment") from error

    def available_ordinals(self) -> list[int]:
        with self._connect() as connection:
            rows = connection.execute(
                "SELECT worker_ordinal FROM assignments WHERE state IN ('assigned','running','result')"
            ).fetchall()
        occupied = {int(row[0]) for row in rows}
        return [ordinal for ordinal in range(3) if ordinal not in occupied]

    def active_assignments(self) -> list[dict[str, Any]]:
        with self._connect() as connection:
            rows = connection.execute(
                "SELECT * FROM assignments WHERE state IN ('assigned','running','result')"
            ).fetchall()
        return [self._record(row) or {} for row in rows]

    def offer(self, ordinal: int) -> dict[str, Any] | None:
        """Return the ordinal's durable assignment, preserving restart identity."""
        now = time.time()
        with self._lock, self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            row = connection.execute(
                """SELECT * FROM assignments WHERE worker_ordinal=?
                AND state IN ('assigned','running') ORDER BY created_at LIMIT 1""",
                (ordinal,),
            ).fetchone()
            if row is not None:
                connection.execute(
                    """UPDATE assignments SET state='running',lease_until=?,last_heartbeat=?,updated_at=?
                    WHERE board=? AND task_id=? AND run_id=?""",
                    (now + self.lease_seconds, now, now, row["board"], row["task_id"], row["run_id"]),
                )
            connection.commit()
        return self._record(row)

    def _matching(self, connection: sqlite3.Connection, envelope: dict[str, Any]) -> sqlite3.Row:
        row = connection.execute(
            "SELECT * FROM assignments WHERE board=? AND task_id=? AND run_id=?",
            (envelope["board"], envelope["task_id"], envelope["run_id"]),
        ).fetchone()
        if row is None:
            raise ProtocolError("assignment is unknown or stale")
        if int(row["worker_ordinal"]) != int(envelope["worker_ordinal"]):
            raise ProtocolError("worker ordinal does not own this assignment")
        if int(row["attempt"]) != int(envelope["attempt"]):
            raise ProtocolError("assignment attempt is stale")
        return row

    def heartbeat(self, envelope: dict[str, Any]) -> tuple[bool, bool]:
        now = time.time()
        delivery_digest = payload_digest(
            {
                name: envelope[name]
                for name in (
                    "kind", "board", "task_id", "run_id", "worker_ordinal",
                    "attempt", "payload_digest",
                )
            }
        )
        with self._lock, self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            duplicate = connection.execute(
                "SELECT digest FROM deliveries WHERE delivery_id=?",
                (envelope["delivery_id"],),
            ).fetchone()
            row = self._matching(connection, envelope)
            if row["state"] not in {"assigned", "running"}:
                raise ProtocolError("assignment is no longer running")
            if row["state"] == "running" and float(row["lease_until"]) < now:
                raise ProtocolError("assignment lease expired")
            if duplicate and duplicate[0] != delivery_digest:
                raise ProtocolError("delivery identifier was reused")
            if not duplicate:
                connection.execute(
                    "INSERT INTO deliveries VALUES (?,?,?,?)",
                    (envelope["delivery_id"], "heartbeat", delivery_digest, now),
                )
            connection.execute(
                "UPDATE assignments SET state='running',lease_until=?,last_heartbeat=?,updated_at=? WHERE board=? AND task_id=? AND run_id=?",
                (now + self.lease_seconds, now, now, envelope["board"], envelope["task_id"], envelope["run_id"]),
            )
            connection.commit()
        return True, bool(duplicate)

    def accept_result(self, envelope: dict[str, Any]) -> tuple[dict[str, Any], bool]:
        now = time.time()
        with self._lock, self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            row = self._matching(connection, envelope)
            digest = envelope["payload_digest"]
            if row["result_digest"]:
                if row["result_digest"] != digest:
                    raise ProtocolError("conflicting result for completed delivery")
                connection.rollback()
                return self._record(row) or {}, True
            if row["state"] not in {"assigned", "running"}:
                raise ProtocolError("assignment cannot accept a result")
            if row["state"] == "running" and float(row["lease_until"]) < now:
                raise ProtocolError("assignment lease expired")
            connection.execute(
                """UPDATE assignments SET state='result',result_digest=?,result_json=?,updated_at=?
                WHERE board=? AND task_id=? AND run_id=?""",
                (digest, canonical_json(envelope["payload"]).decode(), now,
                 envelope["board"], envelope["task_id"], envelope["run_id"]),
            )
            connection.commit()
            row = connection.execute(
                "SELECT * FROM assignments WHERE board=? AND task_id=? AND run_id=?",
                (envelope["board"], envelope["task_id"], envelope["run_id"]),
            ).fetchone()
        return self._record(row) or {}, False

    def pending_results(self) -> list[dict[str, Any]]:
        with self._connect() as connection:
            rows = connection.execute("SELECT * FROM assignments WHERE state='result' ORDER BY updated_at").fetchall()
        return [self._record(row) or {} for row in rows]

    def expire_leases(
        self, *, now: float | None = None, max_attempts: int = 3
    ) -> list[dict[str, Any]]:
        """Fence expired attempts and re-offer or terminally release their ordinals."""
        current = time.time() if now is None else float(now)
        maximum = max(1, min(int(max_attempts), 10))
        changed: list[dict[str, Any]] = []
        with self._lock, self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            rows = connection.execute(
                """SELECT * FROM assignments WHERE state='running'
                AND lease_until > 0 AND lease_until < ? ORDER BY updated_at""",
                (current,),
            ).fetchall()
            for row in rows:
                if int(row["attempt"]) >= maximum:
                    state, attempt = "lease_failed", int(row["attempt"])
                else:
                    state, attempt = "assigned", int(row["attempt"]) + 1
                connection.execute(
                    """UPDATE assignments SET state=?,attempt=?,lease_until=0,
                    last_heartbeat=0,updated_at=? WHERE board=? AND task_id=? AND run_id=?
                    AND attempt=? AND state='running'""",
                    (
                        state, attempt, current, row["board"], row["task_id"],
                        row["run_id"], row["attempt"],
                    ),
                )
                updated = connection.execute(
                    "SELECT * FROM assignments WHERE board=? AND task_id=? AND run_id=?",
                    (row["board"], row["task_id"], row["run_id"]),
                ).fetchone()
                if updated is not None:
                    changed.append(self._record(updated) or {})
            connection.commit()
        return changed

    def finalize(self, binding: dict[str, Any], state: str) -> None:
        if state not in {"finalized", "stale"}:
            raise ProtocolError("invalid terminal assignment state")
        with self._lock, self._connect() as connection:
            connection.execute(
                "UPDATE assignments SET state=?,updated_at=? WHERE board=? AND task_id=? AND run_id=? AND worker_ordinal=? AND attempt=?",
                (state, time.time(), *(binding[name] for name in ("board", "task_id", "run_id", "worker_ordinal", "attempt"))),
            )

    def garbage_collect(self, retention_seconds: int) -> int:
        cutoff = time.time() - max(3600, retention_seconds)
        with self._lock, self._connect() as connection:
            cursor = connection.execute(
                "DELETE FROM assignments WHERE state IN ('finalized','stale') AND updated_at < ?",
                (cutoff,),
            )
            connection.execute("DELETE FROM deliveries WHERE received_at < ?", (cutoff,))
            return int(cursor.rowcount)