from __future__ import annotations

import json
import math
import re
import socket
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Final, Protocol
from uuid import uuid4

from k1link.simulation.contracts import ControlProfile
from k1link.simulation.provider_contract import (
    SimulationProviderContractError,
    SimulationProviderProfile,
)

REQUEST_SCHEMA: Final = "missioncore.simulation-worker-request/v1"
RESPONSE_SCHEMA: Final = "missioncore.simulation-worker-response/v1"
STATUS_SCHEMA: Final = "missioncore.simulation-worker-status/v2"
VEHICLE_STATE_SCHEMA: Final = "missioncore.vehicle-state/v1"
COMMAND_ACCEPTANCE_SCHEMA: Final = "missioncore.command-acceptance/v2"
MAX_MESSAGE_BYTES: Final = 64 * 1024
OPERATIONS: Final = frozenset({"status", "live", "start", "command", "stop"})
SAFE_ID_PATTERN: Final = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")


class SimulationWorkerGatewayError(RuntimeError):
    """The worker transport or response violated the private gateway contract."""


class SimulationWorkerUnavailableError(SimulationWorkerGatewayError):
    """The configured worker cannot currently be reached."""


class SimulationWorkerRejectedError(SimulationWorkerGatewayError):
    """The worker safely rejected a valid gateway request."""


class PolygonWorkerGateway(Protocol):
    def status(self) -> dict[str, Any]: ...

    def live(self) -> dict[str, Any]: ...

    def start(
        self,
        *,
        run_id: str,
        mission_core_commit: str,
        idempotency_key: str,
    ) -> dict[str, Any]: ...

    def stop(
        self,
        *,
        run_id: str,
        idempotency_key: str,
    ) -> dict[str, Any]: ...

    def command(
        self,
        *,
        run_id: str,
        speed_mps: float,
        steering_normalized: float,
        idempotency_key: str,
    ) -> dict[str, Any]: ...


class UnixSocketWorkerGateway:
    """Bounded request/response transport between Mission Core and one local worker."""

    def __init__(
        self,
        socket_path: Path,
        *,
        request_timeout_seconds: float = 3.0,
        lifecycle_timeout_seconds: float = 150.0,
    ) -> None:
        path = socket_path.expanduser()
        if not path.is_absolute():
            raise ValueError("worker socket path must be absolute")
        if request_timeout_seconds <= 0 or lifecycle_timeout_seconds <= 0:
            raise ValueError("worker gateway timeouts must be positive")
        self.socket_path = path
        self.request_timeout_seconds = request_timeout_seconds
        self.lifecycle_timeout_seconds = lifecycle_timeout_seconds

    def status(self) -> dict[str, Any]:
        result = self._request("status", {}, timeout_seconds=self.request_timeout_seconds)
        _validate_status(result)
        return result

    def live(self) -> dict[str, Any]:
        result = self._request("live", {}, timeout_seconds=self.request_timeout_seconds)
        _validate_vehicle_state(result)
        return result

    def start(
        self,
        *,
        run_id: str,
        mission_core_commit: str,
        idempotency_key: str,
    ) -> dict[str, Any]:
        result = self._request(
            "start",
            {
                "run_id": run_id,
                "mission_core_commit": mission_core_commit,
                "idempotency_key": idempotency_key,
            },
            timeout_seconds=self.lifecycle_timeout_seconds,
        )
        _validate_status(result)
        return result

    def stop(
        self,
        *,
        run_id: str,
        idempotency_key: str,
    ) -> dict[str, Any]:
        result = self._request(
            "stop",
            {
                "run_id": run_id,
                "idempotency_key": idempotency_key,
            },
            timeout_seconds=self.lifecycle_timeout_seconds,
        )
        _validate_status(result)
        return result

    def command(
        self,
        *,
        run_id: str,
        speed_mps: float,
        steering_normalized: float,
        idempotency_key: str,
    ) -> dict[str, Any]:
        result = self._request(
            "command",
            {
                "run_id": run_id,
                "speed_mps": speed_mps,
                "steering_normalized": steering_normalized,
                "idempotency_key": idempotency_key,
            },
            timeout_seconds=self.request_timeout_seconds,
        )
        _validate_command_acceptance(result)
        return result

    def _request(
        self,
        operation: str,
        payload: Mapping[str, object],
        *,
        timeout_seconds: float,
    ) -> dict[str, Any]:
        if operation not in OPERATIONS:
            raise ValueError("unsupported worker operation")
        request_id = uuid4().hex
        encoded = (
            json.dumps(
                {
                    "schema_version": REQUEST_SCHEMA,
                    "request_id": request_id,
                    "operation": operation,
                    "payload": dict(payload),
                },
                ensure_ascii=True,
                separators=(",", ":"),
            ).encode("utf-8")
            + b"\n"
        )
        if len(encoded) > MAX_MESSAGE_BYTES:
            raise ValueError("worker request exceeds the bounded message size")
        try:
            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
                client.settimeout(timeout_seconds)
                client.connect(str(self.socket_path))
                client.sendall(encoded)
                response_bytes = _read_line(client)
        except (FileNotFoundError, ConnectionError, OSError, TimeoutError) as exc:
            raise SimulationWorkerUnavailableError(
                "Simulation Worker недоступен через локальный шлюз."
            ) from exc
        response = _json_object(response_bytes, "worker response")
        if set(response) != {"schema_version", "request_id", "ok", "result", "error"}:
            raise SimulationWorkerGatewayError("worker response keys do not match v1")
        if response["schema_version"] != RESPONSE_SCHEMA or response["request_id"] != request_id:
            raise SimulationWorkerGatewayError("worker response identity does not match request")
        if not isinstance(response["ok"], bool):
            raise SimulationWorkerGatewayError("worker response ok flag is invalid")
        if response["ok"]:
            if response["error"] is not None or not isinstance(response["result"], dict):
                raise SimulationWorkerGatewayError("successful worker response is malformed")
            return response["result"]
        error = response["error"]
        if response["result"] is not None or not isinstance(error, str) or not error.strip():
            raise SimulationWorkerGatewayError("rejected worker response is malformed")
        raise SimulationWorkerRejectedError(error[:512])


def _read_line(connection: socket.socket) -> bytes:
    chunks: list[bytes] = []
    size = 0
    while True:
        chunk = connection.recv(min(4096, MAX_MESSAGE_BYTES + 1 - size))
        if not chunk:
            raise SimulationWorkerGatewayError("worker closed the gateway without a response")
        newline = chunk.find(b"\n")
        if newline >= 0:
            chunks.append(chunk[:newline])
            size += newline
            if size > MAX_MESSAGE_BYTES:
                raise SimulationWorkerGatewayError("worker response exceeds the size limit")
            return b"".join(chunks)
        chunks.append(chunk)
        size += len(chunk)
        if size > MAX_MESSAGE_BYTES:
            raise SimulationWorkerGatewayError("worker response exceeds the size limit")


def _json_object(value: bytes, label: str) -> dict[str, Any]:
    try:
        document = json.loads(value.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SimulationWorkerGatewayError(f"{label} is not valid UTF-8 JSON") from exc
    if not isinstance(document, dict) or any(not isinstance(key, str) for key in document):
        raise SimulationWorkerGatewayError(f"{label} must be a JSON object")
    return document


def _validate_status(value: Mapping[str, Any]) -> None:
    expected = {
        "schema_version",
        "worker_id",
        "transport",
        "mode",
        "available",
        "control_available",
        "active_run_id",
        "run_state",
        "active_provider_ids",
        "provider_profile",
        "isolation",
        "authority",
    }
    if set(value) != expected or value.get("schema_version") != STATUS_SCHEMA:
        raise SimulationWorkerGatewayError("worker status does not match v2")
    active_provider_ids = value.get("active_provider_ids")
    if (
        not isinstance(value.get("worker_id"), str)
        or value.get("transport") != "unix"
        or value.get("mode") != "simulation"
        or value.get("available") is not True
        or not isinstance(value.get("control_available"), bool)
        or not isinstance(active_provider_ids, list)
        or len(active_provider_ids) > 32
        or any(
            not isinstance(item, str) or not SAFE_ID_PATTERN.fullmatch(item)
            for item in active_provider_ids
        )
        or len(active_provider_ids) != len(set(active_provider_ids))
        or not isinstance(value.get("isolation"), dict)
        or not isinstance(value.get("authority"), dict)
    ):
        raise SimulationWorkerGatewayError("worker status contains invalid values")
    try:
        SimulationProviderProfile.from_dict(value.get("provider_profile"))
    except SimulationProviderContractError as exc:
        raise SimulationWorkerGatewayError("worker provider profile is invalid") from exc
    if value.get("active_run_id") is not None and not isinstance(value["active_run_id"], str):
        raise SimulationWorkerGatewayError("worker active run id is invalid")
    if value.get("run_state") is not None and not isinstance(value["run_state"], str):
        raise SimulationWorkerGatewayError("worker run state is invalid")


def _validate_vehicle_state(value: Mapping[str, Any]) -> None:
    expected = {
        "schema_version",
        "run_id",
        "sequence",
        "observed_at_utc",
        "host_monotonic_ns",
        "sim_time_ns",
        "frame_id",
        "child_frame_id",
        "pose",
        "source",
        "safety",
    }
    if set(value) != expected or value.get("schema_version") != VEHICLE_STATE_SCHEMA:
        raise SimulationWorkerGatewayError("vehicle state does not match v1")
    if (
        not isinstance(value.get("run_id"), str)
        or not isinstance(value.get("sequence"), int)
        or not isinstance(value.get("observed_at_utc"), str)
        or not isinstance(value.get("host_monotonic_ns"), int)
        or not isinstance(value.get("sim_time_ns"), int)
        or value.get("frame_id") != "map_enu"
        or value.get("child_frame_id") != "base_link_flu"
        or not isinstance(value.get("pose"), dict)
        or not isinstance(value.get("source"), dict)
        or not isinstance(value.get("safety"), dict)
    ):
        raise SimulationWorkerGatewayError("vehicle state contains invalid values")
    source = value["source"]
    if (
        set(source) != {"provider", "topic", "signal", "quality"}
        or not isinstance(source.get("provider"), str)
        or not SAFE_ID_PATTERN.fullmatch(source["provider"])
        or not isinstance(source.get("topic"), str)
        or not 1 <= len(source["topic"]) <= 512
        or "\x00" in source["topic"]
        or source.get("signal") != "ground-truth"
        or source.get("quality") != "diagnostic"
    ):
        raise SimulationWorkerGatewayError("vehicle state source is invalid")


def _validate_command_acceptance(value: Mapping[str, Any]) -> None:
    expected = {
        "schema_version",
        "run_id",
        "command_id",
        "sequence",
        "issued_at_sim_ns",
        "valid_until_sim_ns",
        "speed_mps",
        "steering_normalized",
        "authority_scope",
        "delivery",
    }
    if set(value) != expected or value.get("schema_version") != COMMAND_ACCEPTANCE_SCHEMA:
        raise SimulationWorkerGatewayError("command acceptance does not match v2")
    delivery = value.get("delivery")
    numeric_fields = (
        "sequence",
        "issued_at_sim_ns",
        "valid_until_sim_ns",
        "speed_mps",
        "steering_normalized",
    )
    if (
        not isinstance(value.get("run_id"), str)
        or not isinstance(value.get("command_id"), str)
        or any(
            isinstance(value.get(field), bool)
            or not isinstance(value.get(field), int | float)
            or not math.isfinite(value[field])
            for field in numeric_fields
        )
        or value.get("authority_scope") != "virtual-only"
        or not isinstance(delivery, dict)
        or set(delivery)
        != {
            "provider_id",
            "control_profile",
            "accepted",
            "controller_ready",
            "ttl_expired_count",
            "diagnostics",
        }
        or not isinstance(delivery.get("provider_id"), str)
        or not SAFE_ID_PATTERN.fullmatch(delivery["provider_id"])
        or delivery.get("control_profile")
        != ControlProfile.ROVER_SPEED_STEERING_V1.value
        or delivery.get("accepted") is not True
        or not isinstance(delivery.get("controller_ready"), bool)
        or isinstance(delivery.get("ttl_expired_count"), bool)
        or not isinstance(delivery.get("ttl_expired_count"), int)
        or delivery["ttl_expired_count"] < 0
        or not _valid_boolean_diagnostics(delivery.get("diagnostics"))
    ):
        raise SimulationWorkerGatewayError("command acceptance contains invalid values")


def _valid_boolean_diagnostics(value: object) -> bool:
    return (
        isinstance(value, dict)
        and len(value) <= 16
        and all(
            isinstance(key, str)
            and SAFE_ID_PATTERN.fullmatch(key) is not None
            and isinstance(item, bool)
            for key, item in value.items()
        )
    )
