#!/usr/bin/env python3
"""Consume bounded model SSE streams without reusing partial responses.

The module is intentionally provider-neutral.  It understands the event
envelope used by the Responses API, but does not perform network I/O.  A
caller passes the byte chunks it receives and gets a terminal result plus an
atomic checkpoint.  Partial model text is counted for diagnostics and is
never written to the checkpoint or returned as reusable context.
"""

from __future__ import annotations

import codecs
import hashlib
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

from validate_model_budget import ModelBudgetError, validate_provider_response


DEFAULT_MAX_EVENT_BYTES = 256 * 1024
DEFAULT_MAX_RESPONSE_BYTES = 8 * 1024 * 1024
MAX_ERROR_TEXT = 1024
PROFILE_RETRY = {
    "milestone": "contract",
    "contract": "feature",
    "feature": "doc",
    "doc": "doc",
}


class StreamProtocolError(ValueError):
    """Raised when a stream cannot be decoded as bounded SSE."""


@dataclass(frozen=True)
class SSEEvent:
    event: str
    data: str
    event_id: str | None = None


class SSEDecoder:
    """Incrementally decode UTF-8 SSE frames with hard size limits."""

    def __init__(
        self,
        *,
        max_event_bytes: int = DEFAULT_MAX_EVENT_BYTES,
        max_response_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
    ) -> None:
        if max_event_bytes <= 0 or max_response_bytes <= 0:
            raise ValueError("stream byte limits must be positive")
        if max_event_bytes > max_response_bytes:
            raise ValueError("event limit cannot exceed response limit")
        self.max_event_bytes = int(max_event_bytes)
        self.max_response_bytes = int(max_response_bytes)
        self.bytes_received = 0
        self._decoder = codecs.getincrementaldecoder("utf-8")("strict")
        self._line_buffer = ""
        self._data: list[str] = []
        self._event = "message"
        self._event_id: str | None = None
        self._event_bytes = 0

    def _dispatch(self) -> SSEEvent | None:
        if not self._data:
            self._event = "message"
            self._event_id = None
            self._event_bytes = 0
            return None
        event = SSEEvent(self._event, "\n".join(self._data), self._event_id)
        self._data = []
        self._event = "message"
        self._event_id = None
        self._event_bytes = 0
        return event

    def _line(self, line: str) -> SSEEvent | None:
        # A blank line commits one SSE event.  Ignore comments and unknown
        # fields, as required by the SSE framing rules.
        if line == "":
            return self._dispatch()
        self._event_bytes += len(line.encode("utf-8")) + 1
        if self._event_bytes > self.max_event_bytes:
            raise StreamProtocolError("SSE event exceeds byte budget")
        if line.startswith(":"):
            return None
        field, separator, value = line.partition(":")
        if separator and value.startswith(" "):
            value = value[1:]
        if field == "event":
            self._event = value or "message"
        elif field == "data":
            self._data.append(value)
        elif field == "id":
            self._event_id = value
        return None

    def feed(self, chunk: bytes | bytearray | memoryview) -> list[SSEEvent]:
        if not isinstance(chunk, (bytes, bytearray, memoryview)):
            raise TypeError("SSE chunks must be bytes-like")
        raw = bytes(chunk)
        self.bytes_received += len(raw)
        if self.bytes_received > self.max_response_bytes:
            raise StreamProtocolError("SSE response exceeds byte budget")
        try:
            self._line_buffer += self._decoder.decode(raw, final=False)
        except UnicodeDecodeError as exc:
            raise StreamProtocolError("SSE response is not valid UTF-8") from exc
        events: list[SSEEvent] = []
        while True:
            newline = self._line_buffer.find("\n")
            if newline < 0:
                if len(self._line_buffer.encode("utf-8")) > self.max_event_bytes:
                    raise StreamProtocolError("SSE line exceeds byte budget")
                break
            line = self._line_buffer[:newline]
            self._line_buffer = self._line_buffer[newline + 1 :]
            if line.endswith("\r"):
                line = line[:-1]
            event = self._line(line)
            if event is not None:
                events.append(event)
        return events

    def finish(self) -> list[SSEEvent]:
        try:
            self._line_buffer += self._decoder.decode(b"", final=True)
        except UnicodeDecodeError as exc:
            raise StreamProtocolError("SSE response ends with incomplete UTF-8") from exc
        events: list[SSEEvent] = []
        if self._line_buffer:
            line = self._line_buffer
            self._line_buffer = ""
            event = self._line(line.rstrip("\r"))
            if event is not None:
                events.append(event)
        event = self._dispatch()
        if event is not None:
            events.append(event)
        return events


def _json_event(event: SSEEvent) -> dict[str, Any]:
    try:
        value = json.loads(event.data)
    except (TypeError, json.JSONDecodeError) as exc:
        raise StreamProtocolError("SSE data is not a JSON object") from exc
    if not isinstance(value, dict):
        raise StreamProtocolError("SSE data must be a JSON object")
    return value


def _event_type(event: SSEEvent, payload: dict[str, Any]) -> str:
    value = payload.get("type")
    return str(value) if value is not None else event.event


def _nested_response(payload: dict[str, Any]) -> dict[str, Any]:
    response = payload.get("response")
    return response if isinstance(response, dict) else payload


def _error_info(payload: dict[str, Any]) -> tuple[str, str]:
    response = _nested_response(payload)
    error = response.get("error")
    if isinstance(error, dict):
        code = str(error.get("code", "provider_error"))
        message = str(error.get("message", "provider response failed"))
    else:
        code = str(response.get("code", "provider_error"))
        message = str(response.get("message", "provider response failed"))
    return code[:128], message[:MAX_ERROR_TEXT]


def _context_error(code: str, message: str) -> bool:
    text = f"{code} {message}".lower()
    return "context" in text and ("exceed" in text or "length" in text or "window" in text)


def _retry(profile: str, reason: str) -> dict[str, Any]:
    next_profile = PROFILE_RETRY.get(profile, "doc")
    context_overflow = "context" in reason.lower() and (
        "exceed" in reason.lower() or "length" in reason.lower()
    )
    return {
        "action": "split_task" if profile == "doc" and context_overflow else "rebuild_packet",
        "profile": next_profile,
        "reuse_packet": False,
        "drop_optional_evidence": True,
        "reason": reason,
    }


def _checkpoint(path: Path | None, value: dict[str, Any]) -> str | None:
    if path is None:
        return None
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
    temporary = path.with_name(f".{path.name}.tmp")
    try:
        with temporary.open("wb") as handle:
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass
    return str(path)


def _result(
    *,
    context_id: str,
    packet_id: str,
    profile: str,
    status: str,
    decoder: SSEDecoder,
    event_count: int,
    output_bytes: int,
    output_digest: Any,
    response_id: str | None,
    last_event_id: str | None,
    reason_code: str | None = None,
    error: str | None = None,
    usage: dict[str, int] | None = None,
    finish_reason: str | None = None,
    checkpoint_path: Path | None = None,
) -> dict[str, Any]:
    reusable = status == "completed"
    result: dict[str, Any] = {
        "schema": "model-stream-result.v1",
        "context_id": context_id,
        "packet_id": packet_id,
        "status": status,
        "reusable": reusable,
        "event_count": event_count,
        "bytes_received": decoder.bytes_received,
        "output_bytes": output_bytes,
        "output_sha256": output_digest.hexdigest(),
        "response_id": response_id,
        "last_event_id": last_event_id,
        "reason_code": reason_code,
        "error": error,
        "finish_reason": finish_reason,
        "retry": None if reusable else _retry(profile, reason_code or status),
    }
    if usage is not None:
        result["usage"] = usage
    checkpoint = {
        "schema": "model-stream-checkpoint.v1",
        "context_id": context_id,
        "packet_id": packet_id,
        "status": status,
        "reusable": reusable,
        "reason_code": reason_code,
        "event_count": event_count,
        "bytes_received": decoder.bytes_received,
        "output_bytes": output_bytes,
        "output_sha256": output_digest.hexdigest(),
        "response_id": response_id,
        "last_event_id": last_event_id,
        "finish_reason": finish_reason,
        "retry": result["retry"],
    }
    if usage is not None:
        checkpoint["usage"] = usage
    result["checkpoint_path"] = _checkpoint(checkpoint_path, checkpoint)
    return result


def consume_sse(
    chunks: Iterable[bytes],
    *,
    context_id: str,
    packet_id: str,
    profile: str = "feature",
    max_output_tokens: int = 1024,
    max_output_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
    checkpoint_path: Path | None = None,
    max_event_bytes: int = DEFAULT_MAX_EVENT_BYTES,
    max_response_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
) -> dict[str, Any]:
    """Consume a stream and return a terminal, checkpointed result.

    Any disconnect, malformed frame, provider failure, context overflow, or
    output overrun is non-reusable.  The returned retry plan explicitly asks
    the caller to rebuild a smaller packet; no partial output is included.
    """

    decoder = SSEDecoder(
        max_event_bytes=max_event_bytes,
        max_response_bytes=max_response_bytes,
    )
    if max_output_tokens < 0 or max_output_bytes <= 0:
        raise ValueError("output limits must be non-negative/positive")
    event_count = 0
    output_bytes = 0
    output_digest = hashlib.sha256()
    response_id: str | None = None
    last_event_id: str | None = None
    terminal: dict[str, Any] | None = None

    def handle(event: SSEEvent) -> dict[str, Any] | None:
        nonlocal event_count, output_bytes, response_id, last_event_id
        event_count += 1
        last_event_id = event.event_id or last_event_id
        payload = _json_event(event)
        kind = _event_type(event, payload)
        response = _nested_response(payload)
        if response.get("id") is not None:
            response_id = str(response["id"])
        delta = payload.get("delta")
        if isinstance(delta, str):
            encoded = delta.encode("utf-8")
            output_bytes += len(encoded)
            output_digest.update(encoded)
            if output_bytes > max_output_bytes:
                return _result(
                    context_id=context_id, packet_id=packet_id, profile=profile,
                    status="output_budget_exceeded", decoder=decoder,
                    event_count=event_count, output_bytes=output_bytes,
                    output_digest=output_digest, response_id=response_id,
                    last_event_id=last_event_id, reason_code="output_overrun",
                    checkpoint_path=checkpoint_path,
                )
        if kind in {"response.failed", "error"}:
            code, message = _error_info(payload)
            status = "context_budget_exceeded" if _context_error(code, message) else "provider_failed"
            return _result(
                context_id=context_id, packet_id=packet_id, profile=profile,
                status=status, decoder=decoder, event_count=event_count,
                output_bytes=output_bytes, output_digest=output_digest,
                response_id=response_id, last_event_id=last_event_id,
                reason_code=code, error=message, checkpoint_path=checkpoint_path,
            )
        if kind in {"response.completed", "response.incomplete"}:
            try:
                validated = validate_provider_response(
                    payload, max_output_tokens=max_output_tokens
                )
            except ModelBudgetError as exc:
                message = str(exc)
                overrun = "exceeded output token budget" in message
                return _result(
                    context_id=context_id, packet_id=packet_id, profile=profile,
                    status="output_budget_exceeded" if overrun else "protocol_error",
                    decoder=decoder,
                    event_count=event_count, output_bytes=output_bytes,
                    output_digest=output_digest, response_id=response_id,
                    last_event_id=last_event_id,
                    reason_code="output_overrun" if overrun else "invalid_usage",
                    error=message, checkpoint_path=checkpoint_path,
                )
            status = "truncated" if bool(validated.get("truncated")) else "completed"
            return _result(
                context_id=context_id, packet_id=packet_id, profile=profile,
                status=status, decoder=decoder, event_count=event_count,
                output_bytes=output_bytes, output_digest=output_digest,
                response_id=response_id, last_event_id=last_event_id,
                reason_code="length" if status == "truncated" else None,
                usage={key: int(validated[key]) for key in ("input_tokens", "output_tokens", "total_tokens")},
                finish_reason=validated.get("finish_reason"),
                checkpoint_path=checkpoint_path,
            )
        return None

    try:
        for chunk in chunks:
            for event in decoder.feed(chunk):
                terminal = handle(event)
                if terminal is not None:
                    return terminal
        for event in decoder.finish():
            terminal = handle(event)
            if terminal is not None:
                return terminal
    except (ConnectionError, TimeoutError, OSError) as exc:
        return _result(
            context_id=context_id, packet_id=packet_id, profile=profile,
            status="interrupted", decoder=decoder, event_count=event_count,
            output_bytes=output_bytes, output_digest=output_digest,
            response_id=response_id, last_event_id=last_event_id,
            reason_code="stream_disconnected", error=type(exc).__name__,
            checkpoint_path=checkpoint_path,
        )
    except StreamProtocolError as exc:
        return _result(
            context_id=context_id, packet_id=packet_id, profile=profile,
            status="protocol_error", decoder=decoder, event_count=event_count,
            output_bytes=output_bytes, output_digest=output_digest,
            response_id=response_id, last_event_id=last_event_id,
            reason_code="stream_decode_error", error=str(exc),
            checkpoint_path=checkpoint_path,
        )
    except Exception as exc:
        # SDKs use different exception classes for a body that ends before
        # ``response.completed``.  Treat all ordinary iterator exceptions as
        # a disconnect, but keep only the class name in the checkpoint.
        return _result(
            context_id=context_id, packet_id=packet_id, profile=profile,
            status="interrupted", decoder=decoder, event_count=event_count,
            output_bytes=output_bytes, output_digest=output_digest,
            response_id=response_id, last_event_id=last_event_id,
            reason_code="stream_disconnected", error=type(exc).__name__,
            checkpoint_path=checkpoint_path,
        )

    return _result(
        context_id=context_id, packet_id=packet_id, profile=profile,
        status="interrupted", decoder=decoder, event_count=event_count,
        output_bytes=output_bytes, output_digest=output_digest,
        response_id=response_id, last_event_id=last_event_id,
        reason_code="stream_ended_without_completion",
        checkpoint_path=checkpoint_path,
    )


__all__ = ["SSEDecoder", "SSEEvent", "StreamProtocolError", "consume_sse"]