#!/usr/bin/env python3
"""Receive and qualify the Mission Core live shadow stream without K1 access."""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
import select
import socket
import struct
import sys
import time
from collections import Counter
from contextlib import suppress
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, BinaryIO

WIRE_SCHEMA = "missioncore.live-perception-wire/v1"
REPORT_SCHEMA = "missioncore.e12-shadow-transport-report/v1"
WEBSOCKET_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
MAX_HEADER_BYTES = 64 * 1024
MAX_FRAME_BYTES = 3 * 1024 * 1024


class ProbeError(RuntimeError):
    pass


def _read_exact(stream: BinaryIO, size: int) -> bytes:
    parts: list[bytes] = []
    remaining = size
    while remaining:
        chunk = stream.read(remaining)
        if not chunk:
            raise ProbeError("shadow transport ended unexpectedly")
        parts.append(chunk)
        remaining -= len(chunk)
    return b"".join(parts)


def _read_http_headers(stream: BinaryIO) -> tuple[str, dict[str, str]]:
    encoded = bytearray()
    while not encoded.endswith(b"\r\n\r\n"):
        encoded.extend(_read_exact(stream, 1))
        if len(encoded) > MAX_HEADER_BYTES:
            raise ProbeError("websocket response headers exceed the bound")
    try:
        lines = encoded.decode("ascii").split("\r\n")
    except UnicodeDecodeError as exc:
        raise ProbeError("websocket response headers are not ASCII") from exc
    headers: dict[str, str] = {}
    for line in lines[1:]:
        if not line:
            continue
        name, separator, value = line.partition(":")
        if not separator:
            raise ProbeError("malformed websocket response header")
        headers[name.strip().lower()] = value.strip()
    return lines[0], headers


def _send_client_frame(stream: BinaryIO, opcode: int, payload: bytes) -> None:
    if len(payload) > 125:
        raise ProbeError("client control frame exceeds websocket bound")
    mask = os.urandom(4)
    masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
    stream.write(bytes((0x80 | opcode, 0x80 | len(payload))) + mask + masked)
    stream.flush()


def _read_server_frame(stream: BinaryIO) -> tuple[int, bytes]:
    first, second = _read_exact(stream, 2)
    if not first & 0x80:
        raise ProbeError("fragmented websocket frames are not admitted")
    opcode = first & 0x0F
    masked = bool(second & 0x80)
    length = second & 0x7F
    if length == 126:
        length = struct.unpack("!H", _read_exact(stream, 2))[0]
    elif length == 127:
        length = struct.unpack("!Q", _read_exact(stream, 8))[0]
    if length > MAX_FRAME_BYTES:
        raise ProbeError("websocket frame exceeds the shadow transport bound")
    mask = _read_exact(stream, 4) if masked else None
    payload = _read_exact(stream, length)
    if mask is not None:
        payload = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
    return opcode, payload


def _connect(
    host: str,
    port: int,
    path: str,
    token: str,
    timeout_seconds: float,
) -> tuple[socket.socket, BinaryIO]:
    connection = socket.create_connection((host, port), timeout=timeout_seconds)
    connection.settimeout(timeout_seconds)
    stream = connection.makefile("rwb", buffering=0)
    key = base64.b64encode(os.urandom(16)).decode("ascii")
    request = (
        f"GET {path} HTTP/1.1\r\n"
        f"Host: {host}:{port}\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        f"Sec-WebSocket-Key: {key}\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        f"Authorization: Bearer {token}\r\n"
        "\r\n"
    ).encode("ascii")
    stream.write(request)
    stream.flush()
    status, headers = _read_http_headers(stream)
    expected_accept = base64.b64encode(
        hashlib.sha1(f"{key}{WEBSOCKET_GUID}".encode("ascii")).digest()
    ).decode("ascii")
    if not status.startswith("HTTP/1.1 101 "):
        raise ProbeError(f"websocket upgrade failed: {status}")
    if headers.get("sec-websocket-accept") != expected_accept:
        raise ProbeError("websocket accept identity mismatch")
    connection.settimeout(None)
    return connection, stream


def _decode_event(frame: bytes) -> tuple[dict[str, Any], bytes]:
    if len(frame) < 4:
        raise ProbeError("shadow event is truncated")
    header_bytes = struct.unpack("!I", frame[:4])[0]
    if header_bytes < 2 or header_bytes > MAX_HEADER_BYTES:
        raise ProbeError("shadow event header length is invalid")
    boundary = 4 + header_bytes
    if boundary > len(frame):
        raise ProbeError("shadow event header is truncated")
    try:
        header = json.loads(frame[4:boundary])
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ProbeError("shadow event header is invalid") from exc
    if not isinstance(header, dict) or header.get("schema_version") != WIRE_SCHEMA:
        raise ProbeError("shadow event schema is incompatible")
    payload = frame[boundary:]
    if int(header.get("payload_bytes", -1)) != len(payload):
        raise ProbeError("shadow event payload length mismatch")
    if hashlib.sha256(payload).hexdigest() != header.get("payload_sha256"):
        raise ProbeError("shadow event payload digest mismatch")
    if header.get("commands_enabled") is not False:
        raise ProbeError("shadow transport unexpectedly enables commands")
    if header.get("navigation_or_safety_accepted") is not False:
        raise ProbeError("shadow transport unexpectedly claims safety acceptance")
    return header, payload


def _write_report(output_root: Path, report: dict[str, Any]) -> Path:
    root = output_root.expanduser().resolve()
    if os.name == "nt" and root.drive.upper() != "D:":
        raise ProbeError("Windows probe output must stay on drive D")
    root.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    temporary = root / f".{stamp}-e12-shadow-transport.json.tmp"
    destination = root / f"{stamp}-e12-shadow-transport.json"
    encoded = json.dumps(report, indent=2, sort_keys=True).encode("utf-8") + b"\n"
    with temporary.open("xb") as stream:
        stream.write(encoded)
        stream.flush()
        os.fsync(stream.fileno())
    temporary.replace(destination)
    return destination


def run(args: argparse.Namespace) -> dict[str, Any]:
    token = sys.stdin.readline().strip() if args.token_stdin else args.token
    if not token or len(token) < 40:
        raise ProbeError("shadow bearer token is missing or invalid")
    started_epoch_ns = time.time_ns()
    started_monotonic_ns = time.monotonic_ns()
    counts: Counter[str] = Counter()
    payload_bytes: Counter[str] = Counter()
    first_ingress_sequence: int | None = None
    last_ingress_sequence: int | None = None
    ingress_gaps = 0
    session_ids: set[str] = set()
    source_ids: set[str] = set()
    session_end_seen = False

    connection, stream = _connect(
        args.host,
        args.port,
        args.path,
        token,
        args.socket_timeout_seconds,
    )
    try:
        deadline = time.monotonic() + args.max_duration_seconds
        while time.monotonic() < deadline:
            readable, _, _ = select.select(
                [connection],
                [],
                [],
                min(0.5, max(0.0, deadline - time.monotonic())),
            )
            if not readable:
                continue
            opcode, frame = _read_server_frame(stream)
            if opcode == 0x8:
                break
            if opcode == 0x9:
                _send_client_frame(stream, 0xA, frame)
                continue
            if opcode != 0x2:
                raise ProbeError(f"unexpected websocket opcode: {opcode}")
            header, payload = _decode_event(frame)
            sequence = int(header["ingress_sequence"])
            if last_ingress_sequence is not None and sequence > last_ingress_sequence + 1:
                ingress_gaps += sequence - last_ingress_sequence - 1
            if last_ingress_sequence is not None and sequence <= last_ingress_sequence:
                raise ProbeError("shadow ingress sequence is not strictly increasing")
            first_ingress_sequence = first_ingress_sequence or sequence
            last_ingress_sequence = sequence
            modality = str(header["modality"])
            counts[modality] += 1
            payload_bytes[modality] += len(payload)
            session_ids.add(str(header["session_id"]))
            source_ids.add(str(header["source_id"]))
            if modality == "control":
                try:
                    control = json.loads(payload)
                except json.JSONDecodeError as exc:
                    raise ProbeError("control event payload is invalid") from exc
                if control.get("event") == "session-end":
                    session_end_seen = True
                    break
    finally:
        with suppress(Exception):
            _send_client_frame(stream, 0x8, b"")
        stream.close()
        connection.close()

    completed_epoch_ns = time.time_ns()
    report = {
        "schema_version": REPORT_SCHEMA,
        "state": "completed" if session_end_seen else "timed-out",
        "started_at_epoch_ns": started_epoch_ns,
        "completed_at_epoch_ns": completed_epoch_ns,
        "wall_seconds": (time.monotonic_ns() - started_monotonic_ns) / 1_000_000_000,
        "transport": {
            "kind": "ssh-reverse-tunnel-websocket",
            "endpoint": f"{args.host}:{args.port}",
            "worker_has_k1_connection": False,
            "payload_integrity": "sha256-verified-per-event",
            "clock_qualification": "not-qualified-cross-host",
        },
        "authority": {
            "mode": "shadow-diagnostic-only",
            "commands_enabled": False,
            "navigation_or_safety_accepted": False,
        },
        "events": {
            "counts": dict(sorted(counts.items())),
            "payload_bytes": dict(sorted(payload_bytes.items())),
            "first_ingress_sequence": first_ingress_sequence,
            "last_ingress_sequence": last_ingress_sequence,
            "ingress_sequence_gaps": ingress_gaps,
            "session_end_seen": session_end_seen,
            "session_ids": sorted(session_ids),
            "source_ids": sorted(source_ids),
        },
    }
    destination = _write_report(Path(args.output_root), report)
    report["report_path"] = str(destination)
    return report


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="host.docker.internal")
    parser.add_argument("--port", type=int, default=18012)
    parser.add_argument(
        "--path",
        default=(
            "/api/v1/device-plugins/"
            "nodedc.device.xgrids-lixelkity-k1/live-perception-shadow"
        ),
    )
    parser.add_argument("--token", default="")
    parser.add_argument("--token-stdin", action="store_true")
    parser.add_argument("--max-duration-seconds", type=float, default=900.0)
    parser.add_argument("--socket-timeout-seconds", type=float, default=30.0)
    parser.add_argument(
        "--output-root",
        default="D:/NDC_MISSIONCORE/runtime/derived/e12-shadow-transport",
    )
    return parser.parse_args()


if __name__ == "__main__":
    try:
        print(json.dumps(run(_arguments()), indent=2, sort_keys=True))
    except Exception as exc:
        print(f"E12_SHADOW_TRANSPORT_ERROR={type(exc).__name__}:{exc}", file=sys.stderr)
        raise SystemExit(1) from exc
