"""Versioned, bounded JSON frames over a dedicated control socket.

stdout is deliberately NOT a control channel. Any framing or identity error is
fatal to the connection; never try to resynchronize a partially read frame.
"""
from __future__ import annotations

import asyncio
import json
import re
import struct
from dataclasses import dataclass
from typing import Any

VERSION = 1
# The relay protocol is a DISTINCT version plane from the worker protocol:
# a remote launcher multiplexes framed control over a different transport
# and the two must never be conflated.
RELAY_PROTOCOL_VERSION = 1
MAX_FRAME = 1_048_576
MAX_CODE = 262_144
_ID = re.compile(r"[A-Za-z0-9_.-]{1,128}\Z")


class ProtocolError(Exception):
    pass


def identifier(value: Any) -> str:
    if not isinstance(value, str) or not _ID.fullmatch(value):
        raise ProtocolError("invalid protocol identifier")
    return value


def fields(value: Any, required: set[str], optional: set[str] = frozenset()) -> dict:
    if not isinstance(value, dict) or not required <= value.keys() or value.keys() - required - optional:
        raise ProtocolError(f"expected fields {sorted(required)}, optional {sorted(optional)}")
    return value


def text(value: Any, maximum: int = MAX_CODE) -> str:
    if not isinstance(value, str) or len(value.encode("utf-8")) > maximum:
        raise ProtocolError("expected bounded UTF-8 text")
    return value


def _unique(pairs: list[tuple[str, Any]]) -> dict:
    result = {}
    for key, value in pairs:
        if key in result:
            raise ProtocolError(f"duplicate JSON key: {key}")
        result[key] = value
    return result


def _nonfinite(value: str) -> Any:
    raise ProtocolError(f"non-finite JSON number: {value}")


def encode(value: dict) -> bytes:
    body = json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8")
    if not 0 < len(body) <= MAX_FRAME:
        raise ProtocolError("frame exceeds byte limit")
    return struct.pack(">I", len(body)) + body


async def read_frame(reader: asyncio.StreamReader) -> dict:
    size = struct.unpack(">I", await reader.readexactly(4))[0]
    if not 0 < size <= MAX_FRAME:
        raise ProtocolError("invalid frame length")
    raw = await reader.readexactly(size)
    try:
        return json.loads(raw.decode("utf-8"), object_pairs_hook=_unique, parse_constant=_nonfinite)
    except (ValueError, UnicodeError, RecursionError) as error:
        raise ProtocolError("invalid JSON frame") from error


@dataclass(frozen=True)
class BindingToken:
    """The host-issued launch binding a worker echoes in Hello.

    This is a binding token, NOT measured identity: it proves this process
    is the one the host launched (the same evidence class as --token), and
    nothing about the target itself. It deliberately carries no account,
    machine, platform or cwd fields, so it cannot be used where a measured
    identity is required. Verified identity comes from an actual probe on
    the target; if the two disagree, that is a refusal, not a merge.
    """

    target_id: str
    namespace_id: str
    namespace_generation: str

    @classmethod
    def parse(cls, value: Any) -> "BindingToken":
        stamp = validate_target_stamp(value)
        return cls(**stamp)

    def as_stamp(self) -> dict:
        return {
            "target_id": self.target_id,
            "namespace_id": self.namespace_id,
            "namespace_generation": self.namespace_generation,
        }


def measured_runtime_identity() -> dict:
    """Measure THIS process's runtime identity on the machine it runs on.

    This is the endpoint's own probe: platform, arch, interpreter path and
    version. It is the only allowed source for the Hello runtime identity;
    the echoed launch binding is never a substitute for it.
    """
    import platform as _platform
    import sys as _sys

    return {
        "platform": _platform.system() or "unknown",
        "arch": _platform.machine() or "unknown",
        "interpreter_path": _sys.executable,
        "interpreter_version": f"{_sys.version_info.major}.{_sys.version_info.minor}.{_sys.version_info.micro}",
        "capabilities": [],
    }


def validate_target_stamp(stamp: Any) -> dict:
    """Validate and normalize a target/namespace stamp on the wire."""
    fields(stamp, {"target_id", "namespace_id", "namespace_generation"})
    return {
        "target_id": identifier(stamp["target_id"]),
        "namespace_id": identifier(stamp["namespace_id"]),
        "namespace_generation": identifier(stamp["namespace_generation"]),
    }


def validate_worker_identity(identity: Any) -> dict:
    """Validate the runtime/target identity a worker reports in its Hello.

    An incompatible or incomplete identity is a clean typed refusal, never
    a crash and never silently mapped to something acceptable.
    """
    fields(identity, {"protocol_version", "relay_protocol_version", "runtime"},
           {"target"})
    if type(identity["protocol_version"]) is not int:
        raise ProtocolError("identity protocol_version must be an integer")
    if type(identity["relay_protocol_version"]) is not int:
        raise ProtocolError("identity relay_protocol_version must be an integer")
    runtime = fields(identity["runtime"],
                     {"platform", "arch", "interpreter_path", "interpreter_version"},
                     {"environment_fingerprint", "capabilities"})
    for field in ("platform", "arch", "interpreter_path", "interpreter_version"):
        value = runtime[field]
        if not isinstance(value, str) or not value.strip():
            raise ProtocolError(f"runtime identity {field} must be non-empty")
    caps = runtime.get("capabilities", [])
    if not isinstance(caps, list):
        raise ProtocolError("capabilities must be a list")
    for cap in caps:
        if not isinstance(cap, dict) or not isinstance(cap.get("capability"), str):
            raise ProtocolError("malformed capability entry")
        if cap["capability"] == "closures":
            if type(cap.get("dill_available")) is not bool:
                raise ProtocolError("closures capability requires a boolean dill_available")
        elif cap["capability"] == "skill_manifest":
            revision = cap.get("revision")
            if not isinstance(revision, str) or not revision.strip():
                raise ProtocolError("skill_manifest capability requires a revision")
        elif cap["capability"] == "relay_framing":
            version = cap.get("version")
            if type(version) is not int:
                raise ProtocolError("relay_framing capability requires an integer version")
            if version > RELAY_PROTOCOL_VERSION:
                raise ProtocolError("relay framing version is newer than this endpoint supports")
    if identity["protocol_version"] != VERSION:
        raise ProtocolError(
            f"unsupported worker protocol version {identity['protocol_version']}; "
            f"this endpoint speaks version {VERSION}"
        )
    if identity["relay_protocol_version"] != RELAY_PROTOCOL_VERSION:
        raise ProtocolError(
            f"unsupported relay protocol version {identity['relay_protocol_version']}; "
            f"this endpoint speaks relay version {RELAY_PROTOCOL_VERSION}"
        )
    return identity


class Wire:
    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter,
                 session_id: str, generation: str,
                 target: dict | None = None):
        self.reader, self.writer = reader, writer
        self.session_id = identifier(session_id)
        self.generation = identifier(generation)
        self.target = validate_target_stamp(target) if target is not None else None
        self._write_lock = asyncio.Lock()

    async def send(self, message: dict) -> None:
        packet = {"version": VERSION, "session_id": self.session_id,
                  "generation": self.generation, "message": message}
        if self.target is not None:
            packet["target"] = dict(self.target)
        packet = encode(packet)
        async with self._write_lock:
            self.writer.write(packet)
            await self.writer.drain()

    async def receive(self) -> dict:
        packet = fields(await read_frame(self.reader),
                        {"version", "session_id", "generation", "message"},
                        {"target"})
        if type(packet["version"]) is not int or packet["version"] != VERSION:
            raise ProtocolError("unsupported protocol version")
        if packet["session_id"] != self.session_id or packet["generation"] != self.generation:
            raise ProtocolError("stale or foreign worker identity")
        # Target/namespace mismatch joins session/generation mismatch in the
        # fatal class: a local frame claiming a stamp is a forged identity,
        # a remote frame missing its stamp cannot be routed.
        if self.target is None:
            if "target" in packet:
                raise ProtocolError("local worker envelope must not claim a remote target stamp")
        else:
            if "target" not in packet:
                raise ProtocolError("remote worker envelope is missing its target stamp")
            actual = validate_target_stamp(packet["target"])
            if actual != self.target:
                raise ProtocolError(
                    f"foreign target stamp: expected {self.target}, got {actual}"
                )
        message = packet["message"]
        if not isinstance(message, dict) or not isinstance(message.get("kind"), str):
            raise ProtocolError("missing message kind")
        return message