"""NDJSON bridge transport, framing, and protocol state machine (Revision 1)."""

from __future__ import annotations

import asyncio
import enum
import json
import re
from typing import Any, Dict, Optional, Set

from backend.errors import AppError

PROTOCOL_VERSION = 1
ID_REGEX = re.compile(r"^[A-Za-z0-9_-]{1,128}$")

# Wire limits (excluding trailing newline)
ENVELOPE_ALLOWANCE = 64 * 1024  # 64 KiB
MAX_TOOL_REQUEST_PAYLOAD = 16 * 1024 * 1024  # 16 MiB
MAX_MODEL_REQUEST_PAYLOAD = 128 * 1024 * 1024  # 128 MiB
MAX_COLLECTION_START_FRAME = 6 * 128 * 1024 * 1024 + ENVELOPE_ALLOWANCE
GLOBAL_MAX_FRAME_BYTES = MAX_COLLECTION_START_FRAME

MAX_INFLIGHT_REMOTE_CALLS = 4


class BridgeState(enum.Enum):
    INIT = "INIT"
    RUNNING = "RUNNING"
    COLLECTING = "COLLECTING"
    COMPLETE = "COMPLETE"
    FAILED = "FAILED"


class NDJSONProtocolError(AppError):
    def __init__(self, message: str):
        super().__init__("execution_failed", message, status_code=500)


class BridgeTransport:
    """Async NDJSON reader and writer over attached container streams."""

    def __init__(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        max_frame_bytes: int = GLOBAL_MAX_FRAME_BYTES,
    ):
        self.reader = reader
        self.writer = writer
        self.max_frame_bytes = max_frame_bytes
        self._buffer = bytearray()
        self.state = BridgeState.INIT
        self._seen_ids: Set[str] = set()
        self._inflight_calls: Set[str] = set()  # set of request IDs awaiting response
        self._b_counter = 0

    def _next_backend_id(self) -> str:
        self._b_counter += 1
        return f"b_{self._b_counter}"

    def validate_id(self, msg_id: str) -> None:
        if not isinstance(msg_id, str) or not ID_REGEX.match(msg_id):
            raise NDJSONProtocolError(f"Invalid message ID: '{msg_id}'")
        if msg_id in self._seen_ids:
            raise NDJSONProtocolError(f"Duplicate message ID: '{msg_id}'")
        self._seen_ids.add(msg_id)

    async def read_message_bounded(self) -> Dict[str, Any]:
        """Bounded frame read using readline with strict length check and type-specific limits."""
        try:
            if not hasattr(self.reader, "read"):
                line = await self.reader.readline()  # Test stream compatibility
            else:
                while True:
                    newline = self._buffer.find(b"\n")
                    if newline >= 0:
                        if newline > self.max_frame_bytes:
                            raise NDJSONProtocolError("Bridge frame exceeded maximum allowed bytes")
                        line = bytes(self._buffer[:newline + 1])
                        del self._buffer[:newline + 1]
                        break
                    if len(self._buffer) > self.max_frame_bytes:
                        raise NDJSONProtocolError("Bridge frame exceeded maximum allowed bytes")
                    chunk = await self.reader.read(min(64 * 1024, self.max_frame_bytes + 1 - len(self._buffer)))
                    if not chunk:
                        line = bytes(self._buffer)
                        self._buffer.clear()
                        break
                    self._buffer.extend(chunk)
        except NDJSONProtocolError:
            raise
        except Exception as exc:
            raise NDJSONProtocolError("Error reading bridge stream") from exc

        if not line:
            if self.state == BridgeState.COMPLETE:
                return {"type": "eof"}
            raise NDJSONProtocolError(f"Premature EOF in state {self.state.value}")

        if not line.endswith(b"\n"):
            raise NDJSONProtocolError("Bridge frame missing trailing newline")

        raw_frame = line[:-1]
        if len(raw_frame) > self.max_frame_bytes:
            raise NDJSONProtocolError("Bridge frame exceeded maximum allowed bytes")

        try:
            msg = await asyncio.to_thread(json.loads, raw_frame)
        except Exception as e:
            raise NDJSONProtocolError("Malformed JSON in bridge frame") from e
        if not isinstance(msg, dict) or set(msg) != {"v", "type", "id", "payload"}:
            raise NDJSONProtocolError("Invalid bridge envelope fields")
        if type(msg["v"]) is not int or msg["v"] != PROTOCOL_VERSION:
            raise NDJSONProtocolError("Unsupported protocol version")
        msg_type = msg["type"]
        self.validate_id(msg["id"])
        if not msg["id"].startswith("a_"):
            raise NDJSONProtocolError("Agent message IDs must use a_ prefix")
        schemas = {
            "tool_request": {"tool": str, "parameters": dict},
            "model_request": {"messages": list, "tools": list},
            "collection_start": {"markdown": str, "warnings": list},
            "artifact_begin": {"transfer_id": str, "name": str, "size_bytes": int},
            "artifact_chunk": {"transfer_id": str, "index": int, "data_base64": str},
            "artifact_end": {"transfer_id": str, "size_bytes": int, "chunks": int},
            "complete": {"accepted_transfer_count": int},
            "error": {"code": str, "message": str},
        }
        if not isinstance(msg_type, str) or msg_type not in schemas:
            raise NDJSONProtocolError("Unknown bridge type")
        payload = msg["payload"]
        schema = schemas[msg_type]
        if not isinstance(payload, dict) or set(payload) != set(schema) or any(type(payload[k]) is not t for k, t in schema.items()):
            raise NDJSONProtocolError("Invalid bridge payload fields")
        if any(t is int and payload[k] < 0 for k, t in schema.items()):
            raise NDJSONProtocolError("Bridge counts must be nonnegative integers")
        if "transfer_id" in payload and not ID_REGEX.fullmatch(payload["transfer_id"]):
            raise NDJSONProtocolError("Invalid artifact transfer ID")
        payload_limit = MAX_TOOL_REQUEST_PAYLOAD if msg_type == "tool_request" else MAX_MODEL_REQUEST_PAYLOAD if msg_type == "model_request" else None
        if payload_limit is not None:
            payload_bytes = await asyncio.to_thread(lambda: len(json.dumps(payload).encode("utf-8")))
            if payload_bytes > payload_limit or len(raw_frame) > payload_limit + ENVELOPE_ALLOWANCE:
                raise NDJSONProtocolError("Bridge payload exceeds type-specific byte limit")
        elif msg_type == "collection_start":
            metadata = {**msg, "payload": {**payload, "markdown": ""}}
            metadata_bytes = await asyncio.to_thread(lambda: len(json.dumps(metadata).encode("utf-8")))
            if metadata_bytes > ENVELOPE_ALLOWANCE:
                raise NDJSONProtocolError("Collection metadata exceeds envelope allowance")
            if len(payload["markdown"].encode("utf-8")) > 128 * 1024 * 1024:
                raise NDJSONProtocolError("Collection markdown exceeds 128 MiB limit")
        elif len(raw_frame) > ENVELOPE_ALLOWANCE + 6 * 64 * 1024:
            raise NDJSONProtocolError("Bridge control frame exceeds byte limit")
        self._transition_on_received(msg_type, msg)
        return msg

    def _transition_on_received(self, msg_type: str, msg: Dict[str, Any]) -> None:
        if self.state in (BridgeState.COMPLETE, BridgeState.FAILED):
            raise NDJSONProtocolError(f"Message received in terminal state {self.state.value}")

        if self.state == BridgeState.INIT:
            raise NDJSONProtocolError(f"Received '{msg_type}' before 'start' was sent")

        if msg_type == "error":
            self.state = BridgeState.FAILED
            return

        if self.state == BridgeState.RUNNING:
            if msg_type in ("tool_request", "model_request"):
                if len(self._inflight_calls) >= MAX_INFLIGHT_REMOTE_CALLS:
                    raise NDJSONProtocolError("Exceeded maximum in-flight remote calls (4)")
                self._inflight_calls.add(msg["id"])

            elif msg_type == "collection_start":
                if self._inflight_calls:
                    raise NDJSONProtocolError(
                        f"Cannot transition to collection: {len(self._inflight_calls)} remote calls still in flight"
                    )
                self.state = BridgeState.COLLECTING

            else:
                raise NDJSONProtocolError(f"Unexpected message '{msg_type}' in RUNNING state")

        elif self.state == BridgeState.COLLECTING:
            if msg_type in ("tool_request", "model_request"):
                raise NDJSONProtocolError(f"Cannot perform '{msg_type}' in COLLECTING state")
            elif msg_type in ("artifact_begin", "artifact_chunk", "artifact_end"):
                pass  # Handled by runner/artifacts
            elif msg_type == "complete":
                self.state = BridgeState.COMPLETE
            else:
                raise NDJSONProtocolError(f"Unexpected message '{msg_type}' in COLLECTING state")

        elif self.state in (BridgeState.COMPLETE, BridgeState.FAILED):
            raise NDJSONProtocolError(f"Message received in terminal state {self.state.value}")

    async def _send_frame(self, msg: Dict[str, Any]) -> None:
        msg_id = msg.get("id")
        self.validate_id(msg_id)
        raw = await asyncio.to_thread(lambda: json.dumps(msg).encode("utf-8") + b"\n")
        if len(raw) - 1 > GLOBAL_MAX_FRAME_BYTES:
            raise NDJSONProtocolError("Outgoing frame exceeds maximum bytes")
        if msg["type"] in ("tool_response", "model_response"):
            payload_bytes = await asyncio.to_thread(lambda: len(json.dumps(msg["payload"]).encode("utf-8")))
            if payload_bytes > MAX_MODEL_REQUEST_PAYLOAD:
                raise NDJSONProtocolError("Outgoing result exceeds payload byte limit")
        self.writer.write(raw)
        await self.writer.drain()

    async def send_start(
        self,
        prompt: str,
        system_instruction: str,
        remaining_ms: int,
        model_descriptor: Dict[str, Any],
    ) -> str:
        """Send start frame to container."""
        if self.state != BridgeState.INIT:
            raise NDJSONProtocolError(f"Cannot send start in state {self.state.value}")

        msg_id = self._next_backend_id()
        frame = {
            "v": PROTOCOL_VERSION,
            "type": "start",
            "id": msg_id,
            "payload": {
                "prompt": prompt,
                "system_instruction": system_instruction,
                "remaining_ms": remaining_ms,
                "model": model_descriptor,
            },
        }
        await self._send_frame(frame)
        self.state = BridgeState.RUNNING
        return msg_id

    async def send_tool_response(
        self,
        reply_to: str,
        result: Optional[Dict[str, Any]],
        error: Optional[Dict[str, Any]],
    ) -> str:
        if reply_to in self._inflight_calls:
            self._inflight_calls.remove(reply_to)

        msg_id = self._next_backend_id()
        payload = {
            "result": result if error is None else None,
            "error": error if error is not None else None,
        }
        frame = {
            "v": PROTOCOL_VERSION,
            "type": "tool_response",
            "id": msg_id,
            "reply_to": reply_to,
            "payload": payload,
        }
        await self._send_frame(frame)
        return msg_id

    async def send_model_response(
        self,
        reply_to: str,
        result: Optional[Dict[str, Any]],
        error: Optional[Dict[str, Any]],
    ) -> str:
        if reply_to in self._inflight_calls:
            self._inflight_calls.remove(reply_to)

        msg_id = self._next_backend_id()
        payload = {
            "result": result if error is None else None,
            "error": error if error is not None else None,
        }
        frame = {
            "v": PROTOCOL_VERSION,
            "type": "model_response",
            "id": msg_id,
            "reply_to": reply_to,
            "payload": payload,
        }
        await self._send_frame(frame)
        return msg_id

    async def send_collection_ready(self, reply_to: str) -> str:
        msg_id = self._next_backend_id()
        frame = {
            "v": PROTOCOL_VERSION,
            "type": "collection_ready",
            "id": msg_id,
            "reply_to": reply_to,
            "payload": {},
        }
        await self._send_frame(frame)
        return msg_id

    async def send_artifact_ack(
        self,
        reply_to: str,
        transfer_id: str,
        decision: str,  # "accept" | "skip" | "stored"
        warning: Optional[Dict[str, Any]] = None,
    ) -> str:
        msg_id = self._next_backend_id()
        frame = {
            "v": PROTOCOL_VERSION,
            "type": "artifact_ack",
            "id": msg_id,
            "reply_to": reply_to,
            "payload": {
                "transfer_id": transfer_id,
                "decision": decision,
                "warning": warning,
            },
        }
        await self._send_frame(frame)
        return msg_id