"""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.models.user_message_event import UserMessageEvent
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,
    "arbi.user_message": UserMessageEvent,
}


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
    error_msg: 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):
            resp = raw_data.get("response", {})
            error_msg = resp.get("error", {}).get("message", "unknown streaming error")
        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,
        error=error_msg,
    )


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_response = False  # True while we're mid-response-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 response and this is NOT a delta, close the block
        if self._streaming_response and event.event_type != "response.output_text.delta":
            print()  # newline to end the response stream
            self._streaming_response = False

        if isinstance(data, UserMessageEvent):
            print(f"\n  {ts}[user] {data.content}")

        elif 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}")
            if isinstance(data, dict):
                _print_context_inline(data.get("context"))

        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_response:
                    print(f"\n  {ts}[response] ", end="", flush=True)
                    self._streaming_response = 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]")
            if isinstance(data, dict):
                _print_context_inline(data.get("context"))

        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')}")

        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_response:
            print()
            self._streaming_response = False
        print_sse_summary(result)


def _count_tokens(text: str) -> int:
    """Count tokens using tiktoken (with fallback to char estimate)."""
    if not text:
        return 0
    try:
        import tiktoken

        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except Exception:
        return len(text) // 4


def _print_context_inline(ctx: dict | None) -> None:
    """Print a compact ledger context line from an SSE event.

    Shows the input token breakdown in logical order:
    system → history → [compaction] → [retrieval] → query → total_input/budget
    Output tokens are shown separately (they're not part of the input budget).
    """
    if ctx is None:
        return
    # Input components — always shown (even when 0)
    parts = [f"{k}={ctx.get(k, 0)}" for k in ("system", "history", "query")]
    # Conditional input components — only shown when non-zero
    for k in ("compaction", "retrieval"):
        val = ctx.get(k, 0)
        if val:
            parts.append(f"{k}={val}")
    total = ctx.get("total_input", "?")
    budget = ctx.get("budget", "?")
    cw = ctx.get("context_window", "?")
    parts.append(f"total_input={total}/{budget}")
    parts.append(f"context_window={cw}")
    # Output tokens shown separately — not part of input budget
    output = ctx.get("output", 0)
    if output:
        parts.append(f"output={output}")
    reasoning = ctx.get("reasoning", 0)
    if reasoning:
        parts.append(f"reasoning={reasoning}")
    # Cost and throughput
    total_cost = ctx.get("total_cost")
    if total_cost is not None:
        parts.append(f"cost=${total_cost:.6f}")
    tps = ctx.get("tokens_per_second")
    if tps is not None:
        parts.append(f"{tps} tok/s")
    trunc = ctx.get("truncated", False)
    if trunc:
        parts.append("TRUNCATED")
    print(f"           {'  '.join(parts)}")
    # System detail breakdown on a second line
    system_detail = ctx.get("system_detail")
    if system_detail:
        detail_parts = [f"{k}={v}" for k, v in system_detail.items()]
        print(f"           system_detail: {'  '.join(detail_parts)}")
    # Aggregate LLM call stats (multi-call agent mode)
    all_llm = ctx.get("all_llm_calls")
    if all_llm:
        llm_parts = [
            f"calls={all_llm.get('count', '?')}",
            f"in={all_llm.get('input_tokens', 0)}",
            f"out={all_llm.get('output_tokens', 0)}",
        ]
        reas = all_llm.get("reasoning_tokens", 0)
        if reas:
            llm_parts.append(f"reas={reas}")
        llm_parts.append(f"total={all_llm.get('total_tokens', 0)}")
        print(f"           all_llm_calls: {'  '.join(llm_parts)}")


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 _build_response_label(result: QueryResult) -> tuple[str, float]:
    """Build the RESPONSE summary label from event stream.

    Returns:
        Tuple of (label_string, last_elapsed_time).
    """
    last_t = 0.0
    last_context: dict | None = None
    for event in result.events:
        evt_t = _get_t(event)
        if evt_t is not None:
            last_t = evt_t
        if event.event_type == "response.completed" and isinstance(event.data, dict):
            last_context = event.data.get("context")

    # Use throughput_tokens from context (total LLM output across all calls)
    # to match tokens_per_second which is throughput_tokens / elapsed_seconds.
    throughput = last_context.get("throughput_tokens") if last_context else None
    tokens = throughput or _count_tokens(result.content)
    parts: list[str] = [f"{tokens} tokens"]

    reasoning = last_context.get("reasoning", 0) if last_context else 0
    if reasoning:
        parts[0] += f" + {reasoning} reasoning"
    elapsed = last_context.get("elapsed_seconds") if last_context else None
    if elapsed is not None:
        parts.append(f"{elapsed}s")
    elif last_t > 0:
        parts.append(f"{last_t:.1f}s")
    tps = last_context.get("tokens_per_second") if last_context else None
    if tps is not None:
        parts.append(f"{tps} tok/s")
    total_cost = last_context.get("total_cost") if last_context else None
    if total_cost is not None:
        parts.append(f"${total_cost:.6f}")

    return " | ".join(parts), last_t


def print_sse_summary(result: QueryResult) -> None:
    """Print answer summary with token counts and timing. Use after live streaming."""
    label, _ = _build_response_label(result)

    print("\n" + "-" * 80)
    print(f"THROUGHPUT ({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 isinstance(data, UserMessageEvent):
            print(f"\n  {ts}[user] {data.content}")

        elif 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}")
            if isinstance(data, dict):
                _print_context_inline(data.get("context"))

        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]")
            if isinstance(data, dict):
                _print_context_inline(data.get("context"))

        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')}")

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

    label, _ = _build_response_label(result)
    answer = result.content

    print("\n" + "-" * 80)
    print(f"THROUGHPUT ({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,
        "stream": True,
    }
    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
    error_msg: str | 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,
    ):
        if resp.status_code >= 400:
            # Read the full error body before raising so the error message includes the server detail
            body = resp.read().decode(errors="replace")
            raise httpx.HTTPStatusError(
                f"Server error '{resp.status_code}' for url '{resp.url}': {body}",
                request=resp.request,
                response=resp,
            )

        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):
                resp_obj = raw_data.get("response", {})
                error_msg = resp_obj.get("error", {}).get("message", "unknown streaming error")
            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,
        error=error_msg,
    )
