"""SSE consumer for the Worldtree Conversation API.

Implements docs/contracts/issues/1.contract.md.
"""

from __future__ import annotations

import json
import re
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any, NamedTuple

import httpx
import httpx_sse

_INT_RE = re.compile(r"^-?\d+$")


class SseId(NamedTuple):
    """Parsed composite SSE wire `id:` per spec §SSE id format."""

    turn_id: int
    seq: int


@dataclass(frozen=True)
class WorkerPhase:
    """SSE event `worker_phase`: agent entered a new processing phase."""

    sse_id: SseId
    phase: str
    turn_id: int


@dataclass(frozen=True)
class Thinking:
    """SSE event `thinking`: incremental thinking content from thinking-enabled models."""

    sse_id: SseId
    content: str


@dataclass(frozen=True)
class Text:
    """SSE event `text`: an incremental response-text delta."""

    sse_id: SseId
    content: str


@dataclass(frozen=True)
class TextBoundary:
    """SSE event `text_boundary`: speakable breakpoint after a `text` event."""

    sse_id: SseId
    kind: str
    char_offset: int
    ts: str


@dataclass(frozen=True)
class ToolStart:
    """SSE event `tool_start`: agent is about to execute a tool."""

    sse_id: SseId
    name: str
    arguments: dict[str, Any]


@dataclass(frozen=True)
class ToolResult:
    """SSE event `tool_result`: a tool call completed."""

    sse_id: SseId
    name: str
    result: Any
    duration_ms: int


@dataclass(frozen=True)
class Done:
    """Terminal SSE event `done`: turn succeeded."""

    sse_id: SseId
    phase: str
    response: str
    model: str
    duration_ms: int
    usage: dict[str, int]


@dataclass(frozen=True)
class Error:
    """Terminal SSE event `error`: turn failed."""

    sse_id: SseId
    phase: str
    message: str
    error_code: str | None


@dataclass(frozen=True)
class Cancelled:
    """Terminal SSE event `cancelled`: turn was cancelled server-side."""

    sse_id: SseId
    phase: str
    turn_id: int
    reason: str | None
    partial_message_id: int | None


@dataclass(frozen=True)
class AwaitingLlmFirstToken:
    """SSE event `awaiting_llm_first_token`: heartbeat during slow first-token.

    Fires at the configured interval (default 5s) during the gap between
    `worker_phase` phase=BuildingPrompt and phase=CallingLLM. Lets clients
    render a live "thinking for Ns…" indicator instead of a frozen line
    during legitimate-slow first-token latency. Stops the moment CallingLLM
    fires (defense-in-depth at three sites); no heartbeat after Cancelled
    or stalled terminal events. Tool round-trip re-entries do NOT re-fire
    heartbeats — INV-201-5 scopes the mechanism to the FIRST gap only.

    `elapsed_ms_since_building_prompt` is server-authoritative
    `time.monotonic()`-based — independent of network latency or clock
    skew, monotonically increasing across the heartbeat sequence.

    See docs/conversation-api-spec.md § awaiting_llm_first_token
    (Worldtree #201, v0.29.0).
    """

    sse_id: SseId
    turn_id: int
    elapsed_ms_since_building_prompt: float


@dataclass(frozen=True)
class AffectUpdate:
    """SSE event `affect_update`: persona-state observability snapshot.

    Two emissions per qualifying turn (persona-enabled agent on non-
    ephemeral session): `status="current"` at turn start carrying the full
    snapshot, `status="scheduled"` after post-turn appraisal kicks off
    (lightweight — `snapshot` is None). Suppressed entirely for persona-
    disabled agents (e.g. `domari`, `muninn`), Tier 3 consumer-defined
    agents (Phase 2.0), and ephemeral sessions.

    Bootstrap reads available via `GET /agents/{agent_id}/persona_state`
    (same `snapshot` shape, requires `persona.read` scope).

    See docs/conversation-api-spec.md § affect_update (Worldtree #204,
    v0.28.0).
    """

    sse_id: SseId
    status: str  # "current" | "scheduled"
    turn_id: int
    snapshot: dict[str, Any] | None  # None when status="scheduled"


Event = (
    WorkerPhase
    | Thinking
    | Text
    | TextBoundary
    | ToolStart
    | ToolResult
    | Done
    | Error
    | Cancelled
    | AffectUpdate
    | AwaitingLlmFirstToken
)


@dataclass(frozen=True)
class AdminEvent:
    """One `/admin/events` envelope (INV-046) — an admin-tier lifecycle event.

    Distinct from the turn-stream `Event` union: this is the process-wide admin
    broadcast stream, not a per-turn stream. `id` is a plain monotonic int
    (resets on restart; heartbeats have id=0). `type` is a dotted namespace
    (session.* / turn.* / key.* / system.*). `data` is a type-specific dict —
    most carry `session_id`; per INV-049 it holds IDs + small metadata only.
    """

    id: int
    type: str
    timestamp: str | None
    data: dict[str, Any]


class MalformedSseId(Exception):
    """Raised when an SSE event's `id:` wire field is missing or non-composite."""

    def __init__(self, raw: str) -> None:
        super().__init__(f"malformed SSE id: {raw[:64]!r}")
        self.raw = raw[:64]


class MalformedSseData(Exception):
    """Raised when an SSE event's `data:` payload is non-empty but fails JSON parsing.

    Sibling to MalformedSseId — both are wire-level protocol violations from the
    server. Empty-data frames are skipped silently per issue #7 INV-001; this
    exception is for non-empty payloads that fail json.loads.
    """

    def __init__(self, raw: str) -> None:
        super().__init__(f"malformed SSE data: {raw[:200]!r}")
        self.raw = raw[:200]


class TurnIdFlip(Exception):
    """Raised when an SSE event's turn_id doesn't match the established/expected turn."""

    def __init__(self, *, established: int, got: int) -> None:
        super().__init__(f"turn_id flip: established={established}, got={got}")
        self.established = established
        self.got = got


class SseConnectFailed(Exception):
    """Raised when the SSE endpoint returned a non-2xx status before the stream opened."""

    def __init__(self, *, status: int, body: bytes) -> None:
        body = body[:1024]
        super().__init__(f"SSE connect failed: status={status}, body={body[:128]!r}")
        self.status = status
        self.body = body


class AgentNotAvailable(SseConnectFailed):
    """Eager 409 from the turn POST (Worldtree v1.0.0b1, #331): the session's
    agent is unavailable, so the turn never launched. Pre-b1 this arrived as a
    200 stream + an in-stream `error` event; b1 surfaces it eagerly. Subclass of
    SseConnectFailed so existing `except SseConnectFailed` handlers still catch
    it — this type just adds the parsed `error_code` + `message`."""

    def __init__(self, *, body: bytes, error_code: str, message: str) -> None:
        super().__init__(status=409, body=body)
        self.error_code = error_code
        self.message = message


class TurnLaunchUnavailable(SseConnectFailed):
    """Eager 503 from the turn POST (Worldtree v1.0.0b1, #331): a transient
    turn-launch failure (loop shutdown / resource exhaustion). RETRYABLE.
    Subclass of SseConnectFailed; adds `error_code`, `message`, `retryable`."""

    retryable = True

    def __init__(self, *, body: bytes, error_code: str, message: str) -> None:
        super().__init__(status=503, body=body)
        self.error_code = error_code
        self.message = message


# Canonical error_codes (Worldtree #331 / v1.0.0b2): 409 -> agent_not_available,
# 503 -> not_ready (retryable; re-pinned from internal_error). Used only as a
# fallback default when the body omits error_code — the real code is surfaced
# verbatim from the {detail:{error_code,message}} envelope.
_EAGER_TURN_FAILURE_CODE = {409: "agent_not_available", 503: "not_ready"}


def _eager_failure_fields(body: bytes, status: int) -> tuple[str, str]:
    """Extract (error_code, message) from an eager turn-launch failure body
    (#331). Accepts the Worldtree `{"detail": {...}}` envelope OR a flat
    `{error_code, message}`; falls back to a status-derived default code and a
    generic message when the body is absent / non-JSON / malformed."""
    try:
        parsed: Any = json.loads(body)
    except (json.JSONDecodeError, ValueError):
        parsed = None
    src: dict[str, Any] = {}
    if isinstance(parsed, dict):
        detail = parsed.get("detail")
        src = detail if isinstance(detail, dict) else parsed
    code = src.get("error_code") or _EAGER_TURN_FAILURE_CODE[status]
    message = src.get("message")
    if not isinstance(message, str):
        message = f"turn launch failed (HTTP {status})"
    return str(code), message


class SseConnectionDropped(Exception):
    """Raised when the HTTP/SSE connection dropped mid-stream."""

    def __init__(self, *, last_seen_sse_id: SseId | None) -> None:
        super().__init__(f"SSE connection dropped; last_seen_sse_id={last_seen_sse_id}")
        self.last_seen_sse_id = last_seen_sse_id


class InvalidLastEventId(Exception):
    """Raised on HTTP 400 from a reconnect request — caller's last_event_id was rejected."""

    def __init__(self, *, raw: str) -> None:
        super().__init__(f"server rejected Last-Event-ID: {raw!r}")
        self.raw = raw


class ResumeTurnFinished(Exception):
    """Raised on HTTP 410 from a reconnect — turn finished while disconnected."""

    def __init__(self, *, turn_id: int) -> None:
        super().__init__(f"resume target finished: turn_id={turn_id}")
        self.turn_id = turn_id


class ResumeBufferExpired(Exception):
    """Raised on HTTP 412 from a reconnect — Last-Event-ID older than replay buffer."""

    def __init__(self, *, turn_id: int, buffered_from_seq: int) -> None:
        super().__init__(
            f"replay buffer expired: turn_id={turn_id}, buffered_from_seq={buffered_from_seq}"
        )
        self.turn_id = turn_id
        self.buffered_from_seq = buffered_from_seq


class CancelTurnNotFound(Exception):
    """Raised on HTTP 404 from cancel — turn finished before cancel arrived (race)."""

    def __init__(self, *, turn_id: int) -> None:
        super().__init__(f"cancel target not found: turn_id={turn_id}")
        self.turn_id = turn_id


class CancelAlreadyCompleted(Exception):
    """Raised on HTTP 409 from cancel — turn already terminated server-side (race)."""

    def __init__(self, *, turn_id: int) -> None:
        super().__init__(f"cancel: turn already completed: turn_id={turn_id}")
        self.turn_id = turn_id


class CancelFailed(Exception):
    """Raised on unexpected cancel response status."""

    def __init__(self, *, status: int, body: bytes) -> None:
        body = body[:1024]
        super().__init__(f"cancel failed: status={status}, body={body[:128]!r}")
        self.status = status
        self.body = body


@dataclass(frozen=True)
class CancelResult:
    """Response envelope from POST /sessions/{id}/turns/{turn_id}/cancel."""

    turn_id: int
    cancelled: bool
    reason: str | None
    partial_message_id: int | None


def _envelope_for_type(body: dict[str, Any], sse_id: SseId) -> Event:
    """Dispatch a parsed JSON body to its typed Event variant."""
    t = body["type"]
    if t == "text":
        return Text(sse_id=sse_id, content=body["content"])
    if t == "worker_phase":
        return WorkerPhase(sse_id=sse_id, phase=body["phase"], turn_id=body["turn_id"])
    if t == "thinking":
        return Thinking(sse_id=sse_id, content=body["content"])
    if t == "text_boundary":
        return TextBoundary(
            sse_id=sse_id,
            kind=body["kind"],
            char_offset=body["char_offset"],
            ts=body["ts"],
        )
    if t == "tool_start":
        return ToolStart(sse_id=sse_id, name=body["name"], arguments=body["arguments"])
    if t == "tool_result":
        return ToolResult(
            sse_id=sse_id,
            name=body["name"],
            result=body["result"],
            duration_ms=body["duration_ms"],
        )
    if t == "done":
        return Done(
            sse_id=sse_id,
            phase=body["phase"],
            response=body["response"],
            model=body["model"],
            duration_ms=body["duration_ms"],
            usage=body["usage"],
        )
    if t == "error":
        return Error(
            sse_id=sse_id,
            phase=body.get("phase", "failed"),
            message=body.get("message", ""),
            error_code=body.get("error_code"),
        )
    if t == "cancelled":
        return Cancelled(
            sse_id=sse_id,
            phase=body["phase"],
            turn_id=body["turn_id"],
            reason=body.get("reason"),
            partial_message_id=body.get("partial_message_id"),
        )
    if t == "awaiting_llm_first_token":
        # Worldtree #201 / v0.29.0: top-level heartbeat during BuildingPrompt
        # → CallingLLM gap. Lets clients render live elapsed-time indicators
        # instead of frozen lines on legitimate-slow first-token latency.
        return AwaitingLlmFirstToken(
            sse_id=sse_id,
            turn_id=body["turn_id"],
            elapsed_ms_since_building_prompt=body["elapsed_ms_since_building_prompt"],
        )
    if t == "affect_update":
        # Worldtree #204 / v0.28.0: persona-state observability event.
        # status="current" carries full snapshot at turn start;
        # status="scheduled" omits snapshot (lightweight post-appraisal-
        # kickoff notification).
        return AffectUpdate(
            sse_id=sse_id,
            status=body["status"],
            turn_id=body["turn_id"],
            snapshot=body.get("snapshot"),
        )
    raise ValueError(f"unknown SSE event type: {t!r}")


async def _iter_events(
    event_source: httpx_sse.EventSource,
    *,
    expected_turn_id: int | None,
) -> AsyncIterator[Event]:
    """Apply INV-002 (sse_id present + in range) and INV-003 (turn_id stable) per event.

    `expected_turn_id=None` means "establish from the first event" (stream_turn semantics).
    `expected_turn_id=N` means "every event must match N" (reconnect_turn semantics — the
    first event is already a flip-candidate per INV-003).
    """
    established = expected_turn_id
    last_sse_id: SseId | None = None
    terminal_seen = False
    try:
        async for sse in event_source.aiter_sse():
            # Issue #7 INV-001: empty-data frames are keepalives — skip silently.
            # ORDERING: this branch fires BEFORE _parse_sse_id; an empty-data event
            # with a malformed id is silently swallowed (intentional — a keepalive
            # with a bad id is still a keepalive). Don't reorder.
            if sse.data == "":
                continue
            # v0.8.1: empty-id frames are also treated as keepalives. Worldtree
            # SOMETIMES emits events without an `id:` line (observed mid-stream
            # on the qwen3.6-35-a3b-heretic provider, 2026-05-25). Per the SSE
            # RFC, events without ids are legitimate (they just don't update
            # Last-Event-ID); the previous strict behavior crashed every turn
            # on the offending agent. Treat same as empty-data: skip silently.
            if sse.id == "":
                continue
            try:
                sse_id = _parse_sse_id(sse.id)
            except ValueError as exc:
                raise MalformedSseId(raw=sse.id) from exc
            if established is None:
                established = sse_id.turn_id
            elif sse_id.turn_id != established:
                raise TurnIdFlip(established=established, got=sse_id.turn_id)
            try:
                body = json.loads(sse.data)
            except json.JSONDecodeError as exc:
                raise MalformedSseData(raw=sse.data) from exc
            event = _envelope_for_type(body, sse_id=sse_id)
            yield event
            last_sse_id = sse_id
            if isinstance(event, (Done, Error, Cancelled)):
                terminal_seen = True
                return
    except (httpx.ReadError, httpx.RemoteProtocolError, httpx.ReadTimeout) as exc:
        # ReadTimeout covers idle gaps that exceed httpx's read timeout — the SSE
        # stream went quiet long enough for httpx to give up. Treat the same as a
        # raw read error: surface as SseConnectionDropped so the caller can decide
        # whether to reconnect_turn. (Callers SHOULD configure a long-or-disabled
        # read timeout on their AsyncClient for SSE; this is defense in depth.)
        raise SseConnectionDropped(last_seen_sse_id=last_sse_id) from exc
    if not terminal_seen:
        # Clean EOF before terminal event — INV-001 says stream MUST NOT end
        # without exactly one Done/Error/Cancelled. Surface as connection drop;
        # caller may reconnect_turn if it holds last_sse_id.
        raise SseConnectionDropped(last_seen_sse_id=last_sse_id)


async def stream_turn(
    client: httpx.AsyncClient, session_id: str, content: str
) -> AsyncIterator[Event]:
    """POST a message and yield typed Events. See contract FN stream_turn."""
    assert client is not None
    assert session_id and isinstance(session_id, str)
    assert content and isinstance(content, str)

    async with httpx_sse.aconnect_sse(
        client,
        "POST",
        f"/sessions/{session_id}/messages",
        json={"content": content},
    ) as event_source:
        # Worldtree v1.0.0b1 (#331): turn-launch failures arrive EAGERLY as a
        # status before any stream — 409 agent_not_available (pre-b1 this was a
        # 200 + in-stream `error` event), 503 a transient retryable launch
        # failure. Surface them as typed SseConnectFailed subclasses carrying
        # error_code; request-level non-2xx (404 session_not_found, etc.) stay
        # generic SseConnectFailed.
        status = event_source.response.status_code
        if status in (409, 503):
            body = await event_source.response.aread()
            code, message = _eager_failure_fields(body, status)
            if status == 409:
                raise AgentNotAvailable(body=body, error_code=code, message=message)
            raise TurnLaunchUnavailable(body=body, error_code=code, message=message)
        try:
            event_source.response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            body = await exc.response.aread()
            raise SseConnectFailed(status=exc.response.status_code, body=body) from exc
        async for event in _iter_events(event_source, expected_turn_id=None):
            yield event


async def reconnect_turn(
    client: httpx.AsyncClient,
    session_id: str,
    content: str,
    last_event_id: str,
) -> AsyncIterator[Event]:
    """Re-POST with Last-Event-ID to resume. See contract FN reconnect_turn."""
    assert client is not None
    assert session_id and isinstance(session_id, str)
    assert isinstance(content, str)
    expected = _parse_sse_id(last_event_id)

    async with httpx_sse.aconnect_sse(
        client,
        "POST",
        f"/sessions/{session_id}/messages",
        json={"content": content},
        headers={"Last-Event-ID": last_event_id},
    ) as event_source:
        status = event_source.response.status_code
        if status != 200:
            body_bytes = await event_source.response.aread()
            try:
                body = json.loads(body_bytes)
            except json.JSONDecodeError:
                body = {}
            if status == 400:
                raise InvalidLastEventId(raw=last_event_id)
            if status == 410:
                raise ResumeTurnFinished(turn_id=body.get("turn_id", expected.turn_id))
            if status == 412:
                raise ResumeBufferExpired(
                    turn_id=body.get("turn_id", expected.turn_id),
                    buffered_from_seq=body.get("buffered_from_seq", 0),
                )
            raise SseConnectFailed(status=status, body=body_bytes)
        async for event in _iter_events(event_source, expected_turn_id=expected.turn_id):
            yield event


async def stream_turn_resilient(
    client: httpx.AsyncClient,
    session_id: str,
    content: str,
    *,
    max_reconnects: int = 5,
) -> AsyncIterator[Event]:
    """Resume-orchestration wrapper over stream_turn + reconnect_turn.

    Yields ONE continuous Event stream; on `SseConnectionDropped` (mid-stream
    drop or clean EOF before a terminal), resumes from the last-seen `sse_id`
    via `reconnect_turn`, up to `max_reconnects` times, until a terminal
    Done/Error/Cancelled arrives. The single shared surface presenters consume
    for resilient streaming (design-brief §8b: "share the consumer, branch the
    presenter"). Cross-process resume stays deferred to v2 (§8d): `last_seen`
    lives only in this generator's frame. See contract FN stream_turn_resilient
    (amendment 2026-06-30).
    """
    assert client is not None
    assert session_id and isinstance(session_id, str)
    assert content and isinstance(content, str)
    assert isinstance(max_reconnects, int) and max_reconnects >= 0

    last_seen: SseId | None = None
    reconnects = 0
    gen = stream_turn(client, session_id, content)
    while True:
        try:
            async for event in gen:
                last_seen = event.sse_id
                yield event
            return  # generator completed cleanly → terminal event reached (INV-001)
        except SseConnectionDropped as drop:
            # Prefer the id we tracked from a yielded event; fall back to the one
            # the drop carries (covers a drop on the very first frame). Non-drop
            # reconnect failures (412/410/400/flip) are NOT caught here — they
            # propagate per the contract's "surface, not recover" policy.
            seen = last_seen or drop.last_seen_sse_id
            if seen is None or reconnects >= max_reconnects:
                raise
            reconnects += 1
            gen = reconnect_turn(
                client,
                session_id,
                content,
                last_event_id=f"{seen.turn_id}:{seen.seq}",
            )


async def stream_admin_events(
    client: httpx.AsyncClient,
    *,
    admin_key: str,
    last_event_id: int | None = None,
) -> AsyncIterator[AdminEvent]:
    """GET /admin/events SSE — the admin-tier lifecycle broadcast stream (INV-046).

    Yields `AdminEvent` envelopes as they arrive. Admin-scoped (admin.events.read):
    the request OVERRIDES Authorization with `admin_key` (distinct from the
    client's default consumer bearer). `last_event_id` sets the `Last-Event-ID`
    header for resume (plain decimal int). Long-lived — iterate until the caller
    stops or the connection ends. Non-200 → SseConnectFailed; a mid-stream drop
    → SseConnectionDropped (caller may reconnect from the last-seen `AdminEvent.id`).
    Malformed frames are skipped (best-effort stream).
    """
    assert client is not None
    assert admin_key and isinstance(admin_key, str)
    headers = {"Authorization": f"Bearer {admin_key}"}
    if last_event_id is not None:
        headers["Last-Event-ID"] = str(last_event_id)
    async with httpx_sse.aconnect_sse(
        client, "GET", "/admin/events", headers=headers
    ) as event_source:
        if event_source.response.status_code != 200:
            body = await event_source.response.aread()
            raise SseConnectFailed(status=event_source.response.status_code, body=body)
        try:
            async for sse in event_source.aiter_sse():
                if sse.data == "":
                    continue
                try:
                    env = json.loads(sse.data)
                except json.JSONDecodeError:
                    continue  # skip a malformed admin frame (best-effort)
                yield AdminEvent(
                    id=env.get("id", 0),
                    type=env["type"],
                    timestamp=env.get("timestamp"),
                    data=env.get("data", {}),
                )
        except (httpx.ReadError, httpx.RemoteProtocolError, httpx.ReadTimeout) as exc:
            raise SseConnectionDropped(last_seen_sse_id=None) from exc


def _parse_sse_id(raw: str) -> SseId:
    """Parse the SSE wire `id:` as composite `{turn_id}:{seq}`. See contract FN _parse_sse_id."""
    assert isinstance(raw, str)
    parts = raw.split(":")
    if len(parts) != 2:
        raise ValueError(f"expected '{{turn_id}}:{{seq}}', got: {raw[:64]!r}")
    turn_id_str, seq_str = parts
    if not _INT_RE.match(turn_id_str) or not _INT_RE.match(seq_str):
        raise ValueError(f"expected '{{turn_id}}:{{seq}}' with decimal ints, got: {raw[:64]!r}")
    turn_id = int(turn_id_str)
    seq = int(seq_str)
    if turn_id < 1 or seq < 1:
        raise ValueError(f"expected both ints >= 1 per spec, got: {raw[:64]!r}")
    return SseId(turn_id=turn_id, seq=seq)


async def cancel_turn(
    client: httpx.AsyncClient,
    session_id: str,
    turn_id: int,
    *,
    persist_partial: bool = False,
) -> CancelResult:
    """POST /sessions/{id}/turns/{turn_id}/cancel. See contract FN cancel_turn."""
    assert client is not None
    assert session_id and isinstance(session_id, str)
    assert isinstance(turn_id, int) and turn_id > 0

    params = {"persist_partial": "true"} if persist_partial else None
    resp = await client.post(
        f"/sessions/{session_id}/turns/{turn_id}/cancel", params=params
    )
    if resp.status_code == 404:
        raise CancelTurnNotFound(turn_id=turn_id)
    if resp.status_code == 409:
        raise CancelAlreadyCompleted(turn_id=turn_id)
    if resp.status_code != 200:
        raise CancelFailed(status=resp.status_code, body=resp.content)
    body = resp.json()
    return CancelResult(
        turn_id=body["turn_id"],
        cancelled=body["cancelled"],
        reason=body.get("reason"),
        partial_message_id=body.get("partial_message_id"),
    )
