"""SSE streaming helpers for the ARBI SDK.

All incoming SSE events are parsed into the autogenerated typed models
(``AgentStepEvent``, ``UserInputRequestEvent``, ``ArtifactEvent``) so that
consumers get full type safety.
"""

from __future__ import annotations

import json
import time
from collections.abc import Callable, Iterator
from typing import Any

from arbi_client.models.agent_step_event import AgentStepEvent
from arbi_client.models.agent_step_event_detail_type_0_item import AgentStepEventDetailType0Item
from arbi_client.models.artifact_event import ArtifactEvent
from arbi_client.models.user_input_request_event import UserInputRequestEvent
from arbi_client.sdk._types import QueryResult, SseEvent, StreamingEvent
from arbi_client.types import Unset


def _dget(item: AgentStepEventDetailType0Item, key: str, default: Any = None) -> Any:
    """Safe get for detail items (which store all data in additional_properties)."""
    return item.additional_properties.get(key, default)


#: Mapping of SSE ``event:`` value → typed model class for typed parsing.
#: Only ARBI extension events use typed models; OpenAI events are handled
#: directly in extraction logic with their nested response shapes.
_SSE_TYPE_DISPATCH: dict[str, type] = {
    "arbi.agent_step": AgentStepEvent,
    "arbi.user_input_request": UserInputRequestEvent,
    "arbi.artifact": ArtifactEvent,
}


def parse_sse_event(event_type: str, data: dict[str, Any]) -> SseEvent:
    """Parse a raw JSON dict into the appropriate typed SSE event model.

    Args:
        event_type: The SSE event name (e.g. ``"token"``, ``"agent_step"``).
        data: The JSON-decoded event data.

    Returns:
        A typed model instance.

    Raises:
        ValueError: If *event_type* is not a known SSE event type.
    """
    model_cls = _SSE_TYPE_DISPATCH.get(event_type)
    if model_cls is None:
        raise ValueError(f"Unknown SSE event type: {event_type!r}")
    return model_cls.from_dict(data)


def parse_sse_line_stream(line_iter: Iterator[str]) -> Iterator[tuple[str, dict | str]]:
    """Parse SSE events from an iterator of lines.

    Yields ``(event_type, data)`` tuples where *data* is a JSON-decoded dict
    (or raw string on JSON parse failure).
    """
    current_event = None
    current_data = None

    for line in line_iter:
        line = line.strip()
        if line.startswith("event:"):
            current_event = line[6:].strip()
        elif line.startswith("data:"):
            current_data = line[5:].strip()
        elif line == "" and current_event and current_data:
            try:
                data = json.loads(current_data)
            except json.JSONDecodeError:
                data = current_data
            yield current_event, data
            current_event = None
            current_data = None


def parse_sse_response(sse_text: str) -> QueryResult:
    """Parse a complete SSE response text into a :class:`QueryResult`.

    Args:
        sse_text: Raw SSE formatted text.

    Returns:
        A :class:`QueryResult` with typed events, content, and metadata.
    """
    events: list[StreamingEvent] = []
    content_parts: list[str] = []
    user_input_events: list[UserInputRequestEvent] = []
    metadata: dict | None = None
    assistant_message_ext_id: str | None = None

    for event_type, raw_data in parse_sse_line_stream(sse_text.split("\n")):
        # Parse ARBI extension events into typed models
        if isinstance(raw_data, dict) and event_type in _SSE_TYPE_DISPATCH:
            data: SseEvent | dict | str = parse_sse_event(event_type, raw_data)
        else:
            data = raw_data

        events.append(StreamingEvent(event_type=event_type, data=data))

        # OpenAI Responses API events
        if event_type == "response.created" and isinstance(raw_data, dict):
            resp = raw_data.get("response", {})
            output = resp.get("output", [])
            if output:
                assistant_message_ext_id = output[0].get("id")
        elif event_type == "response.output_text.delta" and isinstance(raw_data, dict):
            delta = raw_data.get("delta", "")
            if delta:
                content_parts.append(delta)
        elif event_type == "response.completed" and isinstance(raw_data, dict):
            resp = raw_data.get("response", {})
            metadata = resp.get("metadata")
        elif event_type == "response.failed" and isinstance(raw_data, dict):
            pass  # error info is in resp.get("error", {}).get("message")
        elif isinstance(data, UserInputRequestEvent):
            user_input_events.append(data)

    return QueryResult(
        content="".join(content_parts),
        events=events,
        user_input_events=user_input_events,
        metadata=metadata,
        assistant_message_ext_id=assistant_message_ext_id,
    )


def _get_elapsed(event: StreamingEvent) -> str:
    """Format elapsed time for display."""
    if event.elapsed is not None:
        return f"{event.elapsed:>7.1f}s "
    data = event.data
    if hasattr(data, "t") and not isinstance(getattr(data, "t", None), Unset):
        t = data.t
        if t is not None:
            return f"{t:>7.1f}s "
    return "         "


def _get_t(event: StreamingEvent) -> float | None:
    """Extract timing from a StreamingEvent."""
    if event.elapsed is not None:
        return event.elapsed
    data = event.data
    if hasattr(data, "t") and not isinstance(getattr(data, "t", None), Unset):
        return data.t
    return None


class SseEventPrinter:
    """Live per-event SSE printer.  Prints each event as it streams in.

    Usage as an ``on_event`` callback::

        printer = SseEventPrinter()
        run_streaming_query(..., on_event=printer)
        printer.finish(result)          # print answer summary

    Or use the module-level :func:`print_sse_event` / :func:`print_sse_summary`
    helpers directly.
    """

    def __init__(self) -> None:
        self._streaming_tokens = False  # True while we're mid-token-stream

    def __call__(self, event: StreamingEvent) -> None:
        """Print a single SSE event live (called per-event during streaming)."""
        data = event.data
        ts = _get_elapsed(event)

        # If we were streaming tokens and this is NOT a delta, close the block
        if self._streaming_tokens and event.event_type != "response.output_text.delta":
            print()  # newline to end the token stream
            self._streaming_tokens = False

        if event.event_type == "response.created":
            resp = data.get("response", {}) if isinstance(data, dict) else {}
            output = resp.get("output", [])
            msg_id = output[0].get("id", "?") if output else "?"
            print(f"\n  {ts}[response.created] id={msg_id}")

        elif isinstance(data, AgentStepEvent):
            _print_agent_step(ts, data)

        elif event.event_type == "response.output_text.delta":
            delta = data.get("delta", "") if isinstance(data, dict) else ""
            if delta:
                if not self._streaming_tokens:
                    print(f"\n  {ts}[tokens] ", end="", flush=True)
                    self._streaming_tokens = True
                print(delta, end="", flush=True)

        elif isinstance(data, UserInputRequestEvent):
            timeout_s = data.timeout_seconds if not isinstance(data.timeout_seconds, Unset) else "?"
            print(f"\n  {ts}[arbi.user_input_request]")
            print(f"           question: {data.question}")
            print(f"           options:  {data.options}")
            print(f"           timeout:  {timeout_s}s")

        elif isinstance(data, ArtifactEvent):
            mode = data.mode if not isinstance(data.mode, Unset) else "?"
            print(f"\n  {ts}[arbi.artifact] {data.title} ({mode}, {len(data.content)} chars)")

        elif event.event_type == "response.completed":
            print(f"\n  {ts}[response.completed] (terminal event)")

        elif event.event_type == "response.failed":
            resp = data.get("response", {}) if isinstance(data, dict) else {}
            err = resp.get("error", {})
            print(f"\n  {ts}[response.failed] {err.get('message', 'unknown error')}")

        elif event.event_type == "arbi.user_message":
            pass

        else:
            raw_str = json.dumps(data)[:200] if isinstance(data, dict) else str(data)[:200]
            print(f"\n  {ts}[{event.event_type}] {raw_str}")

    def finish(self, result: QueryResult) -> None:
        """Print the answer summary after streaming completes."""
        if self._streaming_tokens:
            print()
            self._streaming_tokens = False
        print_sse_summary(result)


def _print_agent_step(ts: str, data: AgentStepEvent) -> None:
    """Format and print a single agent_step event."""
    idx = data.index
    step = data.step if not isinstance(data.step, Unset) else None
    focus = data.focus if not isinstance(data.focus, Unset) else None
    content = data.content if not isinstance(data.content, Unset) else None
    detail = data.detail if not isinstance(data.detail, Unset) else None

    if step == "planning":
        print(f"\n  {ts}[{idx}] PLANNING STEP")
        if content:
            print("  " + "-" * 40)
            for ln in content.splitlines():
                print(f"    {ln}")
            print("  " + "-" * 40)
    elif step == "tool_progress" and detail:
        for td in detail:
            print(f"\n  {ts}     [{_dget(td, 'tool', '?')}] {_dget(td, 'message', '')}")
    elif step == "evaluation" and detail:
        print(f"\n  {ts}[{idx}] EVALUATION ({len(detail)} learnings)")
        for i, ln in enumerate(detail, 1):
            stmt = _dget(ln, "statement", "")
            print(f"           {i}. {stmt}")
    elif detail and not step:
        tools_summary = []
        for td in detail:
            tool = _dget(td, "tool", "?")
            if tool == "search_documents":
                tools_summary.append(f'search_documents(query="{_dget(td, "query", "")}")')
            elif tool == "get_document_passages":
                tools_summary.append(
                    f"get_document_passages(doc={_dget(td, 'doc_ext_id', '?')}, "
                    f"range={_dget(td, 'from_ref')}-{_dget(td, 'to_ref')})"
                )
            elif tool == "get_table_of_contents":
                tools_summary.append(f"get_table_of_contents(doc_count={_dget(td, 'doc_count')})")
            elif tool == "web_search":
                tools_summary.append(f'web_search(query="{_dget(td, "query", "")[:60]}")')
            elif tool == "read_urls":
                url_list = _dget(td, "urls", [])
                tools_summary.append(f"read_urls({len(url_list)} urls)")
            elif tool == "ask_user":
                tools_summary.append(f'ask_user(question="{_dget(td, "question", "")[:60]}")')
            else:
                tools_summary.append(tool)
        print(f"\n  {ts}[{idx}] TOOL CALL  {', '.join(tools_summary)}")
        if focus:
            print(f"           focus: {focus!r}")
    elif step:
        print(f"\n  {ts}[{idx}] {step.upper()}")
    else:
        print(f"\n  {ts}[{idx}] arbi.agent_step (unknown shape)")


def print_sse_summary(result: QueryResult) -> None:
    """Print just the answer summary (content length, timing). Use after live streaming."""
    last_t = 0.0
    for event in result.events:
        evt_t = _get_t(event)
        if evt_t is not None:
            last_t = evt_t

    answer = result.content
    print("\n" + "-" * 80)
    total_label = f" | total: {last_t:.1f}s" if last_t > 0 else ""
    print(f"ANSWER ({len(answer)} chars{total_label})")
    print("-" * 80)

    if result.user_input_events:
        print(f"User input requests: {len(result.user_input_events)}")
    print("=" * 80)


def print_sse_report(result: QueryResult) -> None:
    """Print a detailed report of all SSE events for diagnostics (batch, after streaming).

    For live output during streaming, use :class:`SseEventPrinter` as an ``on_event``
    callback instead.
    """
    print("\n" + "=" * 80)
    print("SSE EVENT STREAM")
    print("=" * 80)

    for event in result.events:
        ts = _get_elapsed(event)
        data = event.data

        if event.event_type == "response.created":
            resp = data.get("response", {}) if isinstance(data, dict) else {}
            output = resp.get("output", [])
            msg_id = output[0].get("id", "?") if output else "?"
            print(f"\n  {ts}[response.created] id={msg_id}")
        elif isinstance(data, AgentStepEvent):
            _print_agent_step(ts, data)
        elif event.event_type == "response.output_text.delta":
            pass  # deltas are aggregated into content for batch report
        elif isinstance(data, UserInputRequestEvent):
            timeout_s = data.timeout_seconds if not isinstance(data.timeout_seconds, Unset) else "?"
            print(f"\n  {ts}[arbi.user_input_request]")
            print(f"           question: {data.question}")
            print(f"           options:  {data.options}")
            print(f"           timeout:  {timeout_s}s")
        elif isinstance(data, ArtifactEvent):
            mode = data.mode if not isinstance(data.mode, Unset) else "?"
            print(f"\n  {ts}[arbi.artifact] {data.title} ({mode}, {len(data.content)} chars)")
        elif event.event_type == "response.completed":
            print(f"\n  {ts}[response.completed] (terminal event)")
        elif event.event_type == "response.failed":
            resp = data.get("response", {}) if isinstance(data, dict) else {}
            err = resp.get("error", {})
            print(f"\n  {ts}[response.failed] {err.get('message', 'unknown error')}")
        elif event.event_type == "arbi.user_message":
            pass
        else:
            raw_str = json.dumps(data)[:200] if isinstance(data, dict) else str(data)[:200]
            print(f"\n  {ts}[{event.event_type}] {raw_str}")

    last_t = 0.0
    for event in result.events:
        evt_t = _get_t(event)
        if evt_t is not None:
            last_t = evt_t

    answer = result.content
    print("\n" + "-" * 80)
    total_label = f" | total: {last_t:.1f}s" if last_t > 0 else ""
    print(f"ANSWER ({len(answer)} chars{total_label})")
    print("-" * 80)
    if answer:
        print(answer[:2000])
    else:
        print("(no answer content)")

    if result.user_input_events:
        print(f"\nUser input requests: {len(result.user_input_events)}")
    print("=" * 80)


def run_streaming_query(
    base_url: str,
    access_token: str,
    query: str,
    tools: dict[str, Any] | None = None,
    config_ext_id: str | None = None,
    *,
    parent_message_ext_id: str | None = None,
    on_event: Callable[[StreamingEvent], None] | None = None,
    timeout: float = 600.0,
    verify: bool = True,
) -> QueryResult:
    """Run a streaming agent query via httpx.

    Args:
        base_url: API base URL (e.g. ``"https://app.arbi.city"``).
        access_token: Bearer token (workspace-scoped JWT).
        query: User query string.
        tools: Optional tool dict.
        config_ext_id: Optional config external ID.
        on_event: Optional callback invoked with a :class:`StreamingEvent`
            for every SSE event as it arrives.
        timeout: HTTP timeout in seconds.
        verify: Whether to verify SSL certificates.

    Returns:
        A :class:`QueryResult` with typed events, content, user_input_events,
        and assistant_message_ext_id.
    """
    import httpx

    request_data: dict[str, Any] = {
        "input": query,
    }
    if tools is not None:
        request_data["tools"] = tools
    if config_ext_id is not None:
        request_data["model"] = config_ext_id
    if parent_message_ext_id is not None:
        request_data["previous_response_id"] = parent_message_ext_id

    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
    }

    all_events: list[StreamingEvent] = []
    content_parts: list[str] = []
    user_input_events: list[UserInputRequestEvent] = []
    assistant_message_ext_id: str | None = None
    metadata: dict | None = None
    t0 = time.monotonic()

    with (
        httpx.Client(timeout=httpx.Timeout(timeout), verify=verify) as http,
        http.stream("POST", f"{base_url.rstrip('/')}/v1/assistant/query", json=request_data, headers=headers) as resp,
    ):
        resp.raise_for_status()

        for event_type, raw_data in parse_sse_line_stream(resp.iter_lines()):
            elapsed = round(time.monotonic() - t0, 2)

            # Parse ARBI extension events into typed models
            if isinstance(raw_data, dict) and event_type in _SSE_TYPE_DISPATCH:
                data: SseEvent | dict | str = parse_sse_event(event_type, raw_data)
            else:
                data = raw_data

            event = StreamingEvent(event_type=event_type, data=data, elapsed=elapsed)
            all_events.append(event)

            # OpenAI Responses API events
            if event_type == "response.created" and isinstance(raw_data, dict):
                resp_obj = raw_data.get("response", {})
                output = resp_obj.get("output", [])
                if output:
                    assistant_message_ext_id = output[0].get("id")
            elif event_type == "response.output_text.delta" and isinstance(raw_data, dict):
                delta = raw_data.get("delta", "")
                if delta:
                    content_parts.append(delta)
            elif event_type == "response.completed" and isinstance(raw_data, dict):
                resp_obj = raw_data.get("response", {})
                metadata = resp_obj.get("metadata")
            elif event_type == "response.failed" and isinstance(raw_data, dict):
                pass  # error info available via resp.get("error", {}).get("message")
            elif isinstance(data, UserInputRequestEvent):
                user_input_events.append(data)

            if on_event:
                on_event(event)

    return QueryResult(
        content="".join(content_parts),
        events=all_events,
        user_input_events=user_input_events,
        metadata=metadata,
        assistant_message_ext_id=assistant_message_ext_id,
    )
