"""SSE streaming helpers for the talk command.

Extracted from talk_cmd.py to reduce file size.
Contains all SSE parsing, streaming display, and stream consumption logic.

Spec source: talk-streaming-transport.md
"""

# @invar:allow function_size: CLI command implementations with interactive loop
# @invar:allow shell_result: talk streaming bridge propagates stream errors to command exit status

from __future__ import annotations

import asyncio
import logging
import json
from dataclasses import dataclass
from pathlib import Path
import sys
from typing import Callable

import httpx
from rich.console import Console

from anima.shell.cli_commands.talk_markdown import StreamMarkdownRenderer
from anima.shell.cli_commands.talk_display import format_turn_metadata

_LOGGER = logging.getLogger(__name__)


@dataclass
class StreamSessionState:
    """State for SSE streaming session.

    Spec source: talk-streaming-transport.md Non-Streaming Fallback
    Tracks whether streaming is supported and cache for entire session.
    """

    streaming_supported: bool | None = None  # None = not probed yet


@dataclass
class SSEEvent:
    """Parsed SSE event.

    Spec source: talk-streaming-transport.md Event Schema
    """

    event: str
    data: dict[str, object]


# @invar:allow shell_result: Pure SSE parsing helper, no I/O
# @shell_orchestration: SSE parser helper scoped to talk CLI for stream support
def parse_sse_line(line: str) -> tuple[str | None, str | None]:
    """Parse a single SSE line into (field, value) or None.

    >>> parse_sse_line('event: chunk')
    ('event', 'chunk')
    >>> parse_sse_line('data: {"turn": 12}')
    ('data', '{"turn": 12}')
    >>> parse_sse_line('')
    (None, None)
    >>> parse_sse_line('invalid')
    (None, None)
    """
    if not line or ":" not in line:
        return (None, None)

    field, _, value = line.partition(":")
    return (field.strip(), value.strip())


# @invar:allow shell_result: Pure SSE parsing helper, no I/O
# @shell_orchestration: SSE parser helper scoped to talk CLI for stream support
# @shell_complexity: SSE event parsing with state machine for field accumulation
def parse_sse_events(lines: list[str]) -> list[SSEEvent]:
    """Parse SSE lines into structured events.

    >>> events = parse_sse_events(['event: chunk', 'data: {"turn": 1, "chunk": "hello"}', '', 'event: done', 'data: {"turn": 1}'])
    >>> len(events)
    2
    >>> events[0].event
    'chunk'
    >>> events[0].data
    {'turn': 1, 'chunk': 'hello'}
    >>> events[1].event
    'done'
    """
    events: list[SSEEvent] = []
    current_event: str | None = None
    current_data: list[str] = []

    for line in lines:
        field, value = parse_sse_line(line)

        if field is None or value is None:
            # Empty line dispatches event
            if current_event is not None and current_data:
                try:
                    data_str = "\n".join(current_data)
                    data_obj = json.loads(data_str) if data_str else {}
                except json.JSONDecodeError:
                    data_obj = {}

                events.append(SSEEvent(event=current_event, data=data_obj))
                current_event = None
                current_data = []
            continue

        if field == "event":
            current_event = value
        elif field == "data":
            current_data.append(value)

    # Handle last event if no trailing empty line
    if current_event is not None and current_data:
        try:
            data_str = "\n".join(current_data)
            data_obj = json.loads(data_str) if data_str else {}
        except json.JSONDecodeError:
            data_obj = {}
        events.append(SSEEvent(event=current_event, data=data_obj))

    return events


# @invar:allow shell_result: SSE probe with exception handling, bool return for fallback logic
# @invar:allow dead_param: Socket path param retained for future UDS URL customization; currently uses localhost for probe
async def probe_streaming_support(client: httpx.AsyncClient, socket_path: Path) -> bool:
    """Probe whether agent supports SSE streaming.

    Spec source: talk-streaming-transport.md Non-Streaming Fallback
    One-time probe for stream support. Returns False on 404/405/connection failure.

    Stream-safe implementation: Uses streaming request to check headers only,
    then closes immediately without consuming response body. Prevents hang
    on infinite SSE streams.

    Args:
        client: httpx async client configured for socket.
        socket_path: MCP socket path (used to construct UDS URL).

    Returns:
        True if /stream endpoint exists and accepts GET.

    Note:
        Regression test in tests/test_talk_streaming_probe.py proves probe
        returns promptly (< 1s) when /stream exists with infinite body.
    """
    try:
        # Stream-safe probe: check headers only, don't consume body
        # This prevents hanging on infinite SSE streams
        async with client.stream("GET", "http://localhost/stream") as response:
            # Any successful response (status < 400) means streaming is supported
            # Close immediately - we only need to verify endpoint exists
            return response.status_code < 400
    except (httpx.HTTPError, OSError, ConnectionError):
        # Connection failure, 404, 405, etc -> fallback to non-streaming
        return False


# @invar:allow shell_result: SSE chunk renderer, no Result needed
# @shell_complexity: Progressive output with JSON/human format dispatch
def _render_stream_chunk(
    event: SSEEvent,
    agent_id: str,
    json_output: bool,
    console: Console,
    *,
    is_first_chunk: bool = False,
) -> None:
    """Render a stream chunk event to console.

    Spec source: talk-streaming-transport.md Event Schema (chunk event)

    Args:
        event: Parsed SSE event with 'chunk' type.
        agent_id: Agent instance ID for display.
        json_output: If True, render as JSON.
        console: Rich console for output.
        is_first_chunk: If True, clear prompt line before output.
    """
    data = event.data
    chunk_text = str(data.get("chunk", ""))

    if json_output:
        # Use pure json.dumps() for machine-readable output - no ANSI escapes.
        # console.print_json() would emit ANSI codes when force_terminal=True.
        payload = {
            "event": "stream_chunk",
            "agent_id": agent_id,
            "turn": data.get("turn"),
            "chunk": chunk_text,
        }
        console.file.write(json.dumps(payload) + "\n")
    else:
        if is_first_chunk:
            # Clear prompt line before first chunk
            sys.stdout.write("\r\033[K")
            sys.stdout.flush()
            # Add "--- " prefix for visual consistency in TTY mode
            if sys.stdout.isatty():
                sys.stdout.write("--- ")
                sys.stdout.flush()
        # Progressive output (no newline, print inline)
        sys.stdout.write(chunk_text)
        sys.stdout.flush()


# @invar:allow shell_result: SSE done renderer, no Result needed
# @shell_complexity: Terminal event rendering with format dispatch
def _render_stream_done(
    event: SSEEvent, agent_id: str, json_output: bool, console: Console
) -> None:
    """Render a stream done event to console.

    Spec source: talk-streaming-transport.md Event Schema (done event)

    Args:
        event: Parsed SSE event with 'done' type.
        agent_id: Agent instance ID for display.
        json_output: If True, render as JSON.
        console: Rich console for output.
    """
    data = event.data

    if json_output:
        # Use pure json.dumps() for machine-readable output - no ANSI escapes.
        payload = {
            "event": "stream_done",
            "agent_id": agent_id,
            "turn": data.get("turn"),
            "stop_reason": data.get("stop_reason"),
            "final_response_chars": data.get("final_response_chars"),
            "usage": data.get("usage"),
        }
        console.file.write(json.dumps(payload) + "\n")
    else:
        # Newline after chunk stream completes.
        # Don't manually write "--- " — prompt_toolkit is handling the prompt.
        # stdin reader already manages the prompt display.
        sys.stdout.write("\n")
        sys.stdout.flush()

        # Extract turn and usage for metadata display
        turn_raw = data.get("turn")
        usage = data.get("usage")

        # Coerce turn to int | None
        turn: int | None = None
        if turn_raw is not None:
            try:
                turn = int(turn_raw)
            except (TypeError, ValueError):
                turn = None

        # Get total_tokens: prefer usage.total_tokens when present,
        # fallback to sum of prompt_tokens + completion_tokens
        total_tokens = 0
        if usage is not None and isinstance(usage, dict):
            # Priority: usage.total_tokens (authoritative) over computed sum
            total_tokens_raw = usage.get("total_tokens")
            if total_tokens_raw is not None:
                try:
                    total_tokens = int(total_tokens_raw)
                except (TypeError, ValueError):
                    total_tokens = 0
            else:
                # Fallback: sum prompt_tokens and completion_tokens
                try:
                    prompt_tokens = usage.get("prompt_tokens", 0) or 0
                    completion_tokens = usage.get("completion_tokens", 0) or 0
                    total_tokens = int(prompt_tokens) + int(completion_tokens)
                except (TypeError, ValueError):
                    # Invalid usage data - default to 0
                    total_tokens = 0

        # Format metadata (plain text, right-aligned)
        # Get console width for right-alignment
        try:
            console_width = console.width
        except Exception:
            console_width = 80
        metadata = format_turn_metadata(
            turn=turn, total_tokens=total_tokens, console_width=console_width
        )
        if metadata:
            # Apply styling to plain text metadata
            console.print(f"[dim cyan]{metadata}[/dim cyan]")


# @invar:allow shell_result: SSE error renderer, no Result needed
# @shell_complexity: Terminal event rendering with format dispatch
def _render_stream_error(
    event: SSEEvent, agent_id: str, json_output: bool, console: Console
) -> None:
    """Render a stream error event to console.

    Spec source: talk-streaming-transport.md Event Schema (error event)

    Args:
        event: Parsed SSE event with 'error' type.
        agent_id: Agent instance ID for display.
        json_output: If True, render as JSON.
        console: Rich console for output.
    """
    data = event.data
    error_msg = str(data.get("error", "Unknown error"))
    error_code = str(data.get("code", "INTERNAL"))

    if json_output:
        # Use pure json.dumps() for machine-readable output - no ANSI escapes.
        payload = {
            "event": "stream_error",
            "agent_id": agent_id,
            "turn": data.get("turn"),
            "error": error_msg,
            "code": error_code,
        }
        console.file.write(json.dumps(payload) + "\n")
    else:
        sys.stdout.write("\r\033[K")
        sys.stdout.flush()
        console.print(f"[red]Stream error ({error_code}): {error_msg}[/red]")


# @invar:allow shell_result: SSE tool_call renderer, no Result needed
# @shell_complexity: Non-terminal event rendering with format dispatch
def _render_stream_tool_call(
    event: SSEEvent,
    agent_id: str,
    json_output: bool,
    console: Console,
) -> None:
    """Render a stream tool_call event to console.

    Spec source: talk-streaming-transport.md Event Schema (tool_call event)

    Args:
        event: Parsed SSE event with 'tool_call' type.
        agent_id: Agent instance ID for display.
        json_output: If True, render as JSON.
        console: Rich console for output.
    """
    data = event.data
    tool_name = str(data.get("tool_name", ""))
    tool_arg_preview = str(data.get("tool_arg_preview", ""))

    if json_output:
        # Use pure json.dumps() for machine-readable output - no ANSI escapes.
        payload = {
            "event": "tool_call",
            "agent_id": agent_id,
            "turn": data.get("turn"),
            "tool_name": tool_name,
            "tool_arg_preview": tool_arg_preview,
        }
        console.file.write(json.dumps(payload) + "\n")
    else:
        # Clear line, then render tool call with dim yellow style
        sys.stdout.write("\r\033[K")
        sys.stdout.flush()
        console.print(f"  [dim yellow]◆ {tool_name} {tool_arg_preview}[/dim yellow]")


# @invar:allow shell_result: SSE stream consumer task
# @invar:allow dead_param: agent_id and chunk_rendered are consumed by deferred renderer callbacks and stream-path coordination
# @shell_complexity: Async SSE line reader with event dispatch and error handling
async def consume_sse_stream(
    *,
    socket_path: Path,
    agent_id: str,
    json_output: bool,
    console: Console,
    stream_complete: asyncio.Event,
    chunk_rendered: asyncio.Event | None = None,
    render_with_prompt: Callable[[Callable[[], None]], None] | None = None,
    use_markdown_renderer: bool = False,
    persona_name: str = "agent",
) -> tuple[bool, str | None]:
    """Consume SSE stream from agent and render events progressively.

    Spec source: talk-streaming-transport.md Transport and Protocol

    Streaming writes directly to stdout without prompt coordination.
    The first chunk clears the prompt line; done/error restore it.
    This avoids the prompt redisplay-per-chunk problem that causes
    visual corruption during progressive output.

    Args:
        socket_path: MCP socket path (used to construct UDS URL).
        agent_id: Agent instance ID for display.
        json_output: If True, render as JSON.
        console: Rich console for output.
        stream_complete: Event to signal when stream terminates.
        chunk_rendered: Event to signal when a chunk is rendered.
        render_with_prompt: Callback to coordinate with prompt.
        use_markdown_renderer: If True, use StreamMarkdownRenderer for block-aware streaming.
        persona_name: Persona label for stream context.

    Returns:
        (had_error, stop_reason) - had_error is True if error event received.
        stop_reason is from done/error event, None if stream failed.
    """
    had_error = False
    stop_reason: str | None = None
    first_chunk_seen = False
    emit_with_prompt = render_with_prompt or (lambda emit: emit())

    # Create markdown renderer if requested and NOT in JSON output mode.
    # The markdown renderer uses rich.Live which emits ANSI escape codes
    # (e.g., cursor hide) that would pollute JSON output.
    markdown_renderer: StreamMarkdownRenderer | None = None
    if use_markdown_renderer and not json_output:
        _LOGGER.debug("Starting markdown stream renderer for persona=%s", persona_name)
        markdown_renderer = StreamMarkdownRenderer(console)
        markdown_renderer.start()

    def _finish_markdown_renderer() -> None:
        if markdown_renderer is None:
            return
        renderer = markdown_renderer
        try:
            emit_with_prompt(lambda: renderer.finish())
        except RuntimeError:
            # Renderer already finished or not started; safe during cleanup paths.
            pass

    try:
        async with httpx.AsyncClient(
            transport=httpx.AsyncHTTPTransport(uds=str(socket_path)),
            timeout=httpx.Timeout(30.0, read=300.0),  # Long read timeout for stream
        ) as client:

            async def _consume_response_lines() -> tuple[bool, str | None]:
                nonlocal had_error, stop_reason, first_chunk_seen
                async with client.stream("GET", "http://localhost/stream") as response:
                    if response.status_code != 200:
                        stream_complete.set()
                        return (True, None)

                    buffer: list[str] = []
                    async for line in response.aiter_lines():
                        buffer.append(line)
                        if not line:
                            events = parse_sse_events(buffer)
                            for event in events:
                                if event.event == "chunk":
                                    chunk_text = str(event.data.get("chunk", ""))
                                    if markdown_renderer is not None:
                                        # Use StreamMarkdownRenderer for block-aware streaming
                                        emit_with_prompt(
                                            lambda ct=chunk_text: (
                                                markdown_renderer.feed(ct)  # type: ignore[union-attr]
                                            )
                                        )
                                    else:
                                        # Use legacy chunk rendering (JSON output or no markdown renderer)
                                        emit_with_prompt(
                                            lambda e=event: _render_stream_chunk(
                                                e,
                                                agent_id,
                                                json_output,
                                                console,
                                                is_first_chunk=not first_chunk_seen,
                                            )
                                        )
                                    if chunk_rendered is not None:
                                        chunk_rendered.set()
                                    first_chunk_seen = True
                                elif event.event == "done":
                                    _finish_markdown_renderer()
                                    emit_with_prompt(
                                        lambda e=event: _render_stream_done(
                                            e, agent_id, json_output, console
                                        )
                                    )
                                    stop_reason = str(
                                        event.data.get("stop_reason", "complete")
                                    )
                                    stream_complete.set()
                                elif event.event == "error":
                                    _finish_markdown_renderer()
                                    emit_with_prompt(
                                        lambda e=event: _render_stream_error(
                                            e, agent_id, json_output, console
                                        )
                                    )
                                    had_error = True
                                    stop_reason = str(event.data.get("code", "UNKNOWN"))
                                    stream_complete.set()
                                elif event.event == "tool_call":
                                    # tool_call is non-terminal: does not set stream_complete
                                    # and does not mark first_chunk_seen
                                    emit_with_prompt(
                                        lambda e=event: _render_stream_tool_call(
                                            e, agent_id, json_output, console
                                        )
                                    )
                                else:
                                    _LOGGER.warning(
                                        "Unknown SSE event type: %s", event.event
                                    )
                            buffer = []
                return (had_error, stop_reason)

            return await _consume_response_lines()

    except (httpx.HTTPError, OSError, ConnectionError) as e:
        # Stream connection error - signal fallback
        if not json_output:
            console.print(f"[red]Stream connection error: {e}[/red]")
        stream_complete.set()
        return (True, None)
    finally:
        _finish_markdown_renderer()

    return (had_error, stop_reason)
