"""BlockStream — event dispatch state machine that emits semantic blocks.

Consumes the raw event stream from ``session.send()`` and produces
a stream of typed blocks with context. Each block carries a
``BlockContext`` identifying which agent produced it.
"""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING

from ._blocks import (
    AnyBlock,
    BlockContext,
    CompactionBlock,
    ErrorBlock,
    FileBlock,
    NativeToolBlock,
    ReasoningBlock,
    ReasoningChunk,
    ReasoningStartBlock,
    ResponseEndBlock,
    ResponseStartBlock,
    RetryBlock,
    TextChunk,
    TextDone,
    ToolExecution,
    ToolGroup,
    ToolResultBlock,
)
from ._events import (
    CompactionInProgress,
    ErrorEvent,
    MessageDone,
    NativeToolCall,
    OutputFileDone,
    ReasoningDelta,
    ReasoningStarted,
    ReasoningSummaryDelta,
    ResponseCancelled,
    ResponseCompleted,
    ResponseCreated,
    ResponseFailed,
    ResponseIncomplete,
    ResponseInProgress,
    ResponseQueued,
    RetryEvent,
    TextDelta,
    ToolCall,
    ToolResult,
)

if TYPE_CHECKING:
    from ._session import Session


def format_tool_args_brief(name: str, arguments: dict[str, object]) -> str:
    """
    Format tool arguments for inline display next to the ``⏵ <name>``
    line.

    Public so frontends that re-render historical tool calls (e.g. the
    REPL's conversation-resume preview) can produce the same
    ``args_summary`` string the live stream produced when the call
    first ran. Without a single source of truth, "this is what you
    saw originally" diverges from "this is what the renderer
    chooses now" the moment either side changes.

    :param name: Tool name, e.g. ``"Read"``.
    :param arguments: Parsed arguments dict, e.g.
        ``{"file_path": "/x/y.py"}``.
    :returns: One-line summary, e.g. ``"y.py"`` for ``Read`` or the
        tool's primary user-meaningful field. Falls back to a JSON
        encoding of the full dict (truncated to 80 chars) for
        unrecognized tools.
    """
    if not arguments:
        return ""
    _KEYS = {
        "Read": "file_path",
        "Write": "file_path",
        "Edit": "file_path",
        "Bash": "command",
        "Glob": "pattern",
        "Grep": "pattern",
        "web_search": "query",
    }
    key = _KEYS.get(name)
    if key and key in arguments:
        s = str(arguments[key])
        if key == "file_path" and "/" in s:
            s = s.rsplit("/", 1)[-1]
        return s[:80] + "…" if len(s) > 80 else s
    try:
        s = json.dumps(arguments, ensure_ascii=False)
    except (TypeError, ValueError):
        s = str(arguments)
    return s[:80] + "…" if len(s) > 80 else s


def _format_native_label(tool_type: str, data: dict[str, object]) -> str:
    """Format a native tool label."""
    if tool_type == "web_search_call":
        action = data.get("action")
        if isinstance(action, dict):
            at = action.get("type", "")
            if at == "search":
                return f"web search: {str(action.get('query', ''))[:80]}"
            if at == "open_page":
                return f"web open: {str(action.get('url', ''))[:80]}"
        return "web search"
    if tool_type == "mcp_call":
        n = data.get("name", "")
        return f"mcp: {n}" if n else "mcp call"
    return tool_type.replace("_", " ")


def _output_text_from_message_content(content: list[dict[str, object]]) -> str:
    """
    Extract completed assistant text from a message item.

    :param content: Message content blocks from ``MessageDone``.
    :returns: Concatenated ``output_text`` content.
    """
    parts: list[str] = []
    for block in content:
        if block.get("type") != "output_text":
            continue
        text = block.get("text")
        if isinstance(text, str):
            parts.append(text)
    return "".join(parts)


class BlockStream:
    """Consumes a session event stream and emits semantic stream blocks.

    :param text_flush_threshold: Min chars to buffer before flushing
        on a word boundary. Default 30.
    """

    def __init__(self, text_flush_threshold: int = 30) -> None:
        self._flush_threshold = text_flush_threshold

    async def stream(
        self,
        session: Session,
        input: str | list[dict[str, object]],
        *,
        files: list[str] | None = None,
    ) -> AsyncIterator[AnyBlock]:
        """Stream semantic blocks for one turn.

        :param session: The :class:`Session` whose event stream to fold.
            The session's ``send()`` method is invoked internally.
        :param input: User text or a list of content-block dicts,
            e.g. ``"hello"`` or ``[{"type": "input_text", "text": "hi"}]``.
        :param files: Optional list of file paths to upload and attach
            to the turn, e.g. ``["./data.csv"]``.
        :returns: Async iterator of :class:`AnyBlock` values in order.
        """
        in_reasoning = False
        reasoning_text = ""
        summary_text = ""
        # Per-section accumulator for reasoning chunks awaiting a
        # newline / threshold flush. Mirrors ``accumulated`` for text.
        reasoning_accumulated = ""
        # Set when this reasoning section streamed any
        # ``ReasoningChunk``. The final ``ReasoningBlock`` is then
        # suppressed so renderers don't show the same text twice
        # (once live, once as a summary panel).
        reasoning_chunks_emitted = False
        in_text = False
        accumulated = ""
        full_text = ""
        pending_tools: dict[str, ToolExecution] = {}
        # Long-lived metadata for tool calls rendered in this stream.
        # Unlike ``pending_tools``, this survives text deltas so a
        # delayed ``ToolResult`` can still render with the original
        # tool name/arguments instead of being dropped.
        tool_executions_by_call_id: dict[str, ToolExecution] = {}
        # Dedup set: every ``call_id`` we've yielded a
        # ``ToolGroup`` for so far. Survives ``pending_tools``
        # clears (which fire on each ``TextDelta`` to flush
        # ``ToolResultBlock``s for completed tools), so the
        # post-stream action_required event for an MCP tool
        # whose inline observed event already rendered doesn't
        # silently re-render once the intervening text deltas
        # cleared ``pending_tools``. Call_ids are SDK-assigned
        # ``toolu_*`` ids — globally unique within this turn,
        # so the set never grows pathologically.
        seen_call_ids: set[str] = set()
        # Dedup set: every ``call_id`` we've yielded a
        # ``ToolResultBlock`` for. Same survives-clear semantics
        # as ``seen_call_ids`` for the call-line side, but on the
        # result side. Necessary because every action_required
        # tool's result fires TWICE on the SSE stream: once
        # inline from ``_dispatch_action_required`` (the moment
        # the dispatch returns) and once from the
        # ``response.completed`` flush in
        # ``_translate_omnigent_event``. Without this, the result
        # panel renders twice.
        seen_result_call_ids: set[str] = set()
        agent: str | None = None
        turn = 0
        started = False

        def _ctx() -> BlockContext:
            depth = agent.count(".") if agent else 0
            return BlockContext(agent=agent, depth=depth, turn=turn)

        async for event in session.send(input, files=files):
            # ── Response lifecycle ───────────────────
            if isinstance(event, ResponseCreated):
                # Tool calls were already yielded immediately. Emit
                # result-only groups for tools that got output between
                # ResponseCompleted and this ResponseCreated.
                for ex in list(pending_tools.values()):
                    if ex.output is not None:
                        yield ToolResultBlock(
                            name=ex.name,
                            call_id=ex.call_id,
                            agent_name=ex.agent_name,
                            output=ex.output,
                            arguments=ex.arguments,
                            args_summary=ex.args_summary,
                            ctx=_ctx(),
                        )
                pending_tools.clear()
                tool_executions_by_call_id.clear()
                agent = event.response.model
                if not started:
                    started = True
                    yield ResponseStartBlock(
                        model=agent,
                        response_id=event.response.id,
                        ctx=_ctx(),
                    )
                else:
                    turn += 1

            elif isinstance(event, ResponseQueued | ResponseInProgress):
                pass

            # ── Reasoning ────────────────────────────
            elif isinstance(event, ReasoningStarted):
                # Already reasoning = a new summarized-thinking block in the
                # same section. Flush this block's tail + a separator so
                # consecutive thought items don't run together.
                if in_reasoning:
                    yield ReasoningChunk(text=reasoning_accumulated + "\n\n", ctx=_ctx())
                    reasoning_accumulated = ""
                    reasoning_chunks_emitted = True
                    continue
                # Entering reasoning closes open text (symmetric with
                # TextDelta closing reasoning) — else interleaved
                # think→speak→think orphans + concatenates text.
                if in_text:
                    if accumulated:
                        yield TextChunk(text=accumulated, ctx=_ctx())
                        accumulated = ""
                    yield TextDone(
                        full_text=full_text,
                        has_code_blocks="```" in full_text,
                        ctx=_ctx(),
                    )
                    in_text = False
                    full_text = ""
                in_reasoning = True
                reasoning_text = ""
                summary_text = ""
                reasoning_accumulated = ""
                reasoning_chunks_emitted = False
                yield ReasoningStartBlock(ctx=_ctx())

            elif isinstance(event, ReasoningDelta | ReasoningSummaryDelta):
                # An out-of-order delta (no preceding ReasoningStarted)
                # marks an implicit start — Codex's bridged events
                # arrive this way. Emit the start block once so the
                # formatter has its "thinking…" anchor.
                if not in_reasoning:
                    # Close any open text first — same boundary as ReasoningStarted.
                    if in_text:
                        if accumulated:
                            yield TextChunk(text=accumulated, ctx=_ctx())
                            accumulated = ""
                        yield TextDone(
                            full_text=full_text,
                            has_code_blocks="```" in full_text,
                            ctx=_ctx(),
                        )
                        in_text = False
                        full_text = ""
                    in_reasoning = True
                    reasoning_text = ""
                    summary_text = ""
                    reasoning_accumulated = ""
                    reasoning_chunks_emitted = False
                    yield ReasoningStartBlock(ctx=_ctx())
                if isinstance(event, ReasoningDelta):
                    reasoning_text += event.delta
                else:
                    summary_text += event.delta
                # Stream the delta as a ReasoningChunk so the TUI can
                # render mid-flight. Mirrors the TextDelta line/
                # threshold flush so chunks land on natural breaks.
                reasoning_accumulated += event.delta
                while "\n" in reasoning_accumulated:
                    line, reasoning_accumulated = reasoning_accumulated.split("\n", 1)
                    yield ReasoningChunk(text=line + "\n", ctx=_ctx())
                    reasoning_chunks_emitted = True
                if len(reasoning_accumulated) >= self._flush_threshold:
                    last_space = reasoning_accumulated.rfind(" ")
                    if last_space > 0:
                        yield ReasoningChunk(
                            text=reasoning_accumulated[: last_space + 1],
                            ctx=_ctx(),
                        )
                        reasoning_accumulated = reasoning_accumulated[last_space + 1 :]
                        reasoning_chunks_emitted = True

            # ── Text ─────────────────────────────────
            elif isinstance(event, TextDelta):
                if in_reasoning:
                    in_reasoning = False
                    if reasoning_accumulated:
                        yield ReasoningChunk(
                            text=reasoning_accumulated,
                            ctx=_ctx(),
                        )
                        reasoning_chunks_emitted = True
                        reasoning_accumulated = ""
                    if not reasoning_chunks_emitted:
                        yield ReasoningBlock(
                            reasoning_text=reasoning_text,
                            summary_text=summary_text,
                            ctx=_ctx(),
                        )
                # Emit results for tools that completed.
                for ex in list(pending_tools.values()):
                    if ex.output is not None:
                        yield ToolResultBlock(
                            name=ex.name,
                            call_id=ex.call_id,
                            agent_name=ex.agent_name,
                            output=ex.output,
                            arguments=ex.arguments,
                            args_summary=ex.args_summary,
                            ctx=_ctx(),
                        )
                pending_tools.clear()

                in_text = True
                accumulated += event.delta
                full_text += event.delta

                while "\n" in accumulated:
                    line, accumulated = accumulated.split("\n", 1)
                    yield TextChunk(text=line + "\n", ctx=_ctx())

                if len(accumulated) >= self._flush_threshold:
                    last_space = accumulated.rfind(" ")
                    if last_space > 0:
                        yield TextChunk(text=accumulated[: last_space + 1], ctx=_ctx())
                        accumulated = accumulated[last_space + 1 :]

            # ── Tool calls ───────────────────────────
            elif isinstance(event, ToolCall):
                # Dedupe by call_id. Under the claude-sdk harness's
                # MCP path, a tool call surfaces as TWO ToolCall
                # events with correlated call_ids: an inline
                # observed event (status="completed") emitted as
                # the inner SDK parses the tool_use block, and a
                # post-stream action_required event emitted when
                # the SDK invokes the MCP-server handler. The
                # adapter (omnigent/runtime/harnesses/
                # _executor_adapter.py) threads the SDK's
                # tool_use_id through both so they share a
                # call_id; this block keeps the first occurrence
                # (the inline render) and drops the second so the
                # REPL doesn't render ``⏵ tool_name`` twice. See
                # designs/RUN_OMNIGENT_REPL_PARITY.md.
                #
                # Non-MCP paths emit exactly one ToolCall per
                # call_id, so the second-arrival branch never
                # fires for them.
                if event.call_id in seen_call_ids:
                    # Already rendered the ⏵ line for this call_id
                    # (e.g. via the inline observed event from the
                    # harness's content_block_stop, or via the
                    # ``ToolCallInProgress`` event the Omnigent server
                    # emits at action_required arrival). Re-register
                    # in pending_tools so the eventual ``ToolResult``
                    # can pair by call_id — the prior pending entry
                    # may have been cleared by the
                    # ``pending_tools.clear()`` that fires on every
                    # ``TextDelta`` between observed and
                    # action_required (or between action_required
                    # and the post-PATCH function_call_output).
                    execution = tool_executions_by_call_id.get(event.call_id)
                    if execution is None:
                        execution = ToolExecution(
                            name=event.name,
                            arguments=event.arguments,
                            args_summary=format_tool_args_brief(event.name, event.arguments),
                            call_id=event.call_id,
                            agent_name=event.agent_name,
                            executed_by="server",
                        )
                        tool_executions_by_call_id[event.call_id] = execution
                    pending_tools[event.call_id] = execution
                    continue
                seen_call_ids.add(event.call_id)

                if in_reasoning:
                    in_reasoning = False
                    if reasoning_accumulated:
                        yield ReasoningChunk(
                            text=reasoning_accumulated,
                            ctx=_ctx(),
                        )
                        reasoning_chunks_emitted = True
                        reasoning_accumulated = ""
                    if not reasoning_chunks_emitted:
                        yield ReasoningBlock(
                            reasoning_text=reasoning_text,
                            summary_text=summary_text,
                            ctx=_ctx(),
                        )
                if in_text:
                    if accumulated:
                        yield TextChunk(text=accumulated, ctx=_ctx())
                        accumulated = ""
                    yield TextDone(
                        full_text=full_text,
                        has_code_blocks="```" in full_text,
                        ctx=_ctx(),
                    )
                    in_text = False
                    full_text = ""

                execution = ToolExecution(
                    name=event.name,
                    arguments=event.arguments,
                    args_summary=format_tool_args_brief(event.name, event.arguments),
                    call_id=event.call_id,
                    agent_name=event.agent_name,
                    executed_by="server",
                )
                pending_tools[event.call_id] = execution
                tool_executions_by_call_id[event.call_id] = execution
                # Yield immediately so the user sees the tool call
                # before execution. output=None means the formatter
                # shows the call line but no result panel.
                yield ToolGroup(executions=[execution], ctx=_ctx())

            elif isinstance(event, ToolResult):
                if event.call_id in seen_result_call_ids:
                    # Already rendered the result panel. The late
                    # ``response.completed`` flush emission is
                    # redundant; drop it so a subsequent
                    # ``ResponseCreated`` doesn't yield it again
                    # via the pending_tools sweep.
                    continue
                ex = pending_tools.get(event.call_id) or tool_executions_by_call_id.get(
                    event.call_id
                )
                if ex is None:
                    # Result arrived after a TextDelta cleared
                    # ``pending_tools``. If the matching call line
                    # never rendered, we have no name/agent for the
                    # panel — drop the event rather than render
                    # with placeholders.
                    continue
                ex.output = event.output
                ex.executed_by = "client"
                if event.arguments:
                    ex.arguments = event.arguments
                # Yield the result panel IMMEDIATELY so multi-
                # tool turns don't bunch result rendering at
                # end-of-turn. Each dispatch's
                # ``_live_publish`` of the function_call_output
                # arrives here as the tool finishes.
                seen_result_call_ids.add(event.call_id)
                yield ToolResultBlock(
                    name=ex.name,
                    call_id=ex.call_id,
                    agent_name=ex.agent_name,
                    output=event.output,
                    arguments=ex.arguments,
                    args_summary=ex.args_summary,
                    ctx=_ctx(),
                )
                # Drop from ``pending_tools`` so the next
                # ``TextDelta`` / ``ResponseCreated`` sweep
                # doesn't re-yield this result.
                pending_tools.pop(event.call_id, None)

            # ── Native tools ─────────────────────────
            elif isinstance(event, NativeToolCall):
                yield NativeToolBlock(
                    tool_type=event.tool_type,
                    label=_format_native_label(event.tool_type, event.data),
                    data=event.data,
                    ctx=_ctx(),
                )

            # ── Message done ─────────────────────────
            elif isinstance(event, MessageDone):
                if in_reasoning:
                    in_reasoning = False
                    if reasoning_accumulated:
                        yield ReasoningChunk(
                            text=reasoning_accumulated,
                            ctx=_ctx(),
                        )
                        reasoning_chunks_emitted = True
                        reasoning_accumulated = ""
                    if not reasoning_chunks_emitted:
                        yield ReasoningBlock(
                            reasoning_text=reasoning_text,
                            summary_text=summary_text,
                            ctx=_ctx(),
                        )
                if in_text:
                    if accumulated:
                        yield TextChunk(text=accumulated, ctx=_ctx())
                        accumulated = ""
                    yield TextDone(
                        full_text=full_text,
                        has_code_blocks="```" in full_text,
                        ctx=_ctx(),
                    )
                    in_text = False
                    full_text = ""
                else:
                    text = _output_text_from_message_content(event.content)
                    if text:
                        # No prior deltas: synthesize a TextChunk so
                        # TextChunk-only consumers (Session._stream_chunks)
                        # still see the message.
                        yield TextChunk(text=text, ctx=_ctx())
                        yield TextDone(
                            full_text=text,
                            has_code_blocks="```" in text,
                            ctx=_ctx(),
                        )

            # ── Status events ────────────────────────
            elif isinstance(event, CompactionInProgress):
                yield CompactionBlock(ctx=_ctx())

            elif isinstance(event, RetryEvent):
                yield RetryBlock(
                    source=event.source,
                    attempt=event.attempt,
                    max_attempts=event.max_attempts,
                    delay_seconds=event.delay_seconds,
                    ctx=_ctx(),
                )

            elif isinstance(event, ErrorEvent):
                # Pass ``code`` through too — renderers need it as a
                # fallback label when ``message`` is empty (otherwise
                # the error panel shows just ``[llm]`` with no hint
                # as to what went wrong).
                yield ErrorBlock(
                    message=event.error.message,
                    source=event.source,
                    code=event.error.code,
                    ctx=_ctx(),
                )

            elif isinstance(event, OutputFileDone):
                yield FileBlock(file_id=event.file_id, filename=event.filename, ctx=_ctx())

            # ── Terminal events ──────────────────────
            elif isinstance(
                event,
                ResponseCompleted | ResponseFailed | ResponseIncomplete | ResponseCancelled,
            ):
                if in_reasoning:
                    in_reasoning = False
                    if reasoning_accumulated:
                        yield ReasoningChunk(
                            text=reasoning_accumulated,
                            ctx=_ctx(),
                        )
                        reasoning_chunks_emitted = True
                        reasoning_accumulated = ""
                    if not reasoning_chunks_emitted:
                        yield ReasoningBlock(
                            reasoning_text=reasoning_text,
                            summary_text=summary_text,
                            ctx=_ctx(),
                        )
                if in_text:
                    if accumulated:
                        yield TextChunk(text=accumulated, ctx=_ctx())
                        accumulated = ""
                    yield TextDone(
                        full_text=full_text,
                        has_code_blocks="```" in full_text,
                        ctx=_ctx(),
                    )
                    in_text = False
                    full_text = ""

                yield ResponseEndBlock(
                    status=event.response.status,
                    response=event.response,
                    ctx=_ctx(),
                )