"""Transport-agnostic agent streaming engine.

Extracts the stream-parsing loop from ``adapter.py`` so that TUI,
headless CLI, and (future) Tauri sidecar all share one implementation.
"""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from collections.abc import AsyncGenerator, Awaitable, Callable
from datetime import UTC, datetime
from typing import Any

from icarus.session.events import (
    AutoApproveEnabled,
    SessionEvent,
    StatusUpdate,
    TextComplete,
    TextDelta,
    TokenUsage,
    ToolCallStart,
    ToolResult,
    ToolsApproved,
    ToolsRejected,
)
from icarus.utils.file_ops import FileOpTracker
from icarus.utils.image_utils import create_multimodal_content
from icarus.utils.input import ImageTracker, parse_file_mentions
from icarus.utils.tool_display import format_tool_message_content
from langchain.agents.middleware.human_in_the_loop import (
    ApproveDecision,
    HITLRequest,
    HITLResponse,
    RejectDecision,
)
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.types import Command, Interrupt
from pydantic import TypeAdapter

logger = logging.getLogger(__name__)

HITLDecision = ApproveDecision | RejectDecision
_HITL_REQUEST_ADAPTER = TypeAdapter(HITLRequest)


# ---------------------------------------------------------------------------
# Public helpers (also used by app.py for thread-history loading)
# ---------------------------------------------------------------------------

def build_stream_config(
    thread_id: str,
    assistant_id: str | None,
) -> dict[str, Any]:
    """Build the LangGraph stream config dict."""
    metadata: dict[str, str] = {}
    if assistant_id:
        metadata.update(
            {
                "assistant_id": assistant_id,
                "agent_name": assistant_id,
                "updated_at": datetime.now(UTC).isoformat(),
            }
        )
    return {
        "configurable": {"thread_id": thread_id},
        "metadata": metadata,
    }


# ---------------------------------------------------------------------------
# Input pre-processing
# ---------------------------------------------------------------------------

def _preprocess_input(
    user_input: str,
    image_tracker: ImageTracker | None = None,
) -> Any:
    """Parse ``@file`` mentions and build the message content payload."""
    prompt_text, mentioned_files = parse_file_mentions(user_input)
    max_embed_bytes = 256 * 1024

    if mentioned_files:
        context_parts = [prompt_text, "\n\n## Referenced Files\n"]
        for file_path in mentioned_files:
            try:
                file_size = file_path.stat().st_size
                if file_size > max_embed_bytes:
                    size_kb = file_size // 1024
                    context_parts.append(
                        f"\n### {file_path.name}\n"
                        f"Path: `{file_path}`\n"
                        f"Size: {size_kb}KB (too large to embed, "
                        "use read_file tool to view)"
                    )
                else:
                    content = file_path.read_text()
                    context_parts.append(
                        f"\n### {file_path.name}\n"
                        f"Path: `{file_path}`\n```\n{content}\n```"
                    )
            except Exception as e:  # noqa: BLE001
                context_parts.append(
                    f"\n### {file_path.name}\n[Error reading file: {e}]"
                )
        final_input = "\n".join(context_parts)
    else:
        final_input = prompt_text

    images_to_send = image_tracker.get_images() if image_tracker else []
    if images_to_send:
        message_content = create_multimodal_content(final_input, images_to_send)
    else:
        message_content = final_input

    if image_tracker:
        image_tracker.clear()

    return message_content


# ---------------------------------------------------------------------------
# Interrupted-state persistence
# ---------------------------------------------------------------------------

def _build_interrupted_ai_message(
    pending_text: str,
    active_tool_calls: dict[str, tuple[str, dict]],
) -> AIMessage | None:
    """Build an ``AIMessage`` that captures interrupted state for the graph."""
    text = pending_text.strip()
    tool_calls = [
        {"id": tid, "name": name, "args": args}
        for tid, (name, args) in active_tool_calls.items()
    ]
    if not text and not tool_calls:
        return None
    return AIMessage(content=text, tool_calls=tool_calls or [])


def _is_summarization_chunk(metadata: dict | None) -> bool:
    if metadata is None:
        return False
    return metadata.get("lc_source") == "summarization"


# ---------------------------------------------------------------------------
# HITL helpers
# ---------------------------------------------------------------------------

def _mark_file_ops_approved(
    tracker: FileOpTracker,
    action_requests: list[dict],
) -> None:
    for ar in action_requests:
        tool_name = ar.get("name")
        if tool_name in {"write_file", "edit_file"}:
            args = ar.get("args", {})
            if isinstance(args, dict):
                tracker.mark_hitl_approved(tool_name, args)


# ---------------------------------------------------------------------------
# Main streaming generator
# ---------------------------------------------------------------------------

async def stream_turn(  # noqa: C901, PLR0912, PLR0915
    *,
    agent: Any,
    user_input: str,
    thread_id: str,
    assistant_id: str | None = None,
    backend: Any = None,
    auto_approve: bool = False,
    image_tracker: ImageTracker | None = None,
    resolve_hitl: Callable[..., Awaitable[dict]] | None = None,
) -> AsyncGenerator[SessionEvent, None]:
    """Stream a single conversational turn, yielding typed events.

    Args:
        agent: The compiled LangGraph agent.
        user_input: Raw user text (may contain ``@file`` mentions).
        thread_id: LangGraph thread identifier.
        assistant_id: Agent/assistant identifier.
        backend: Backend for ``FileOpTracker`` diff computation.
        auto_approve: Skip HITL approval when ``True``.
        image_tracker: Optional image attachments.
        resolve_hitl: ``async (action_requests, assistant_id) -> decision``
            callback.  Required when *auto_approve* is ``False``.
    """
    message_content = _preprocess_input(user_input, image_tracker)
    config = build_stream_config(thread_id, assistant_id)

    file_op_tracker = FileOpTracker(assistant_id=assistant_id, backend=backend)
    displayed_tool_ids: set[str] = set()
    tool_call_buffers: dict[str | int, dict] = {}
    active_tool_calls: dict[str, tuple[str, dict]] = {}

    pending_text = ""
    current_message_id: str | None = None
    captured_tokens = 0

    yield StatusUpdate("Thinking")

    stream_input: dict | Command = {
        "messages": [{"role": "user", "content": message_content}]
    }

    try:
        while True:
            interrupt_occurred = False
            hitl_response: dict[str, HITLResponse] = {}
            pending_interrupts: dict[str, HITLRequest] = {}

            async for chunk in agent.astream(
                stream_input,
                stream_mode=["messages", "updates"],
                subgraphs=True,
                config=config,
                durability="exit",
            ):
                if not isinstance(chunk, tuple) or len(chunk) != 3:  # noqa: PLR2004
                    continue

                namespace, current_stream_mode, data = chunk
                ns_key = tuple(namespace) if namespace else ()
                is_main_agent = ns_key == ()

                # ---- UPDATES stream (interrupts) ----
                if current_stream_mode == "updates":
                    if not isinstance(data, dict):
                        continue
                    if "__interrupt__" in data:
                        interrupts: list[Interrupt] = data["__interrupt__"]
                        for interrupt_obj in interrupts:
                            validated = _HITL_REQUEST_ADAPTER.validate_python(
                                interrupt_obj.value
                            )
                            pending_interrupts[interrupt_obj.id] = validated
                            interrupt_occurred = True
                    continue

                # ---- MESSAGES stream (main agent only) ----
                if current_stream_mode != "messages" or not is_main_agent:
                    continue

                if not isinstance(data, tuple) or len(data) != 2:  # noqa: PLR2004
                    continue

                message, metadata = data

                if _is_summarization_chunk(metadata):
                    continue

                # -- HumanMessage: flush any pending assistant text --
                if isinstance(message, HumanMessage):
                    if pending_text:
                        yield TextComplete(pending_text, current_message_id or "")
                        pending_text = ""
                        current_message_id = None
                    continue

                # -- ToolMessage: tool result --
                if isinstance(message, ToolMessage):
                    tool_content = format_tool_message_content(message.content)
                    record = file_op_tracker.complete_with_message(message)
                    yield StatusUpdate("Thinking")

                    tool_id = getattr(message, "tool_call_id", None) or ""
                    tool_status = getattr(message, "status", "success")
                    output_str = str(tool_content) if tool_content else ""

                    yield ToolResult(
                        tool_id=tool_id,
                        tool_name=getattr(message, "name", ""),
                        status="success" if tool_status == "success" else "error",
                        output=output_str,
                        diff=record.diff if record else None,
                        diff_path=record.display_path if record else None,
                    )
                    active_tool_calls.pop(tool_id, None)
                    continue

                # -- Token usage (may appear on any AIMessageChunk) --
                if hasattr(message, "usage_metadata"):
                    usage = message.usage_metadata
                    if usage:
                        total = usage.get("total_tokens", 0)
                        if total:
                            captured_tokens = max(captured_tokens, total)
                        else:
                            inp = usage.get("input_tokens", 0)
                            out = usage.get("output_tokens", 0)
                            if inp or out:
                                captured_tokens = max(captured_tokens, inp + out)

                if not hasattr(message, "content_blocks"):
                    continue

                # -- Content blocks (text + tool calls) --
                for block in message.content_blocks:
                    block_type = block.get("type")

                    if block_type == "text":
                        text = block.get("text", "")
                        if text:
                            if current_message_id is None:
                                yield StatusUpdate(None)
                                current_message_id = f"asst-{uuid.uuid4().hex[:8]}"
                            pending_text += text
                            yield TextDelta(text, current_message_id)

                    elif block_type in {"tool_call_chunk", "tool_call"}:
                        chunk_name = block.get("name")
                        chunk_args = block.get("args")
                        chunk_id = block.get("id")
                        chunk_index = block.get("index")

                        buffer_key: str | int
                        if chunk_index is not None:
                            buffer_key = chunk_index
                        elif chunk_id is not None:
                            buffer_key = chunk_id
                        else:
                            buffer_key = f"unknown-{len(tool_call_buffers)}"

                        buffer = tool_call_buffers.setdefault(
                            buffer_key,
                            {"name": None, "id": None, "args": None, "args_parts": []},
                        )

                        if chunk_name:
                            buffer["name"] = chunk_name
                        if chunk_id:
                            buffer["id"] = chunk_id

                        if isinstance(chunk_args, dict):
                            buffer["args"] = chunk_args
                            buffer["args_parts"] = []
                        elif isinstance(chunk_args, str):
                            if chunk_args:
                                parts: list[str] = buffer.setdefault("args_parts", [])
                                if not parts or chunk_args != parts[-1]:
                                    parts.append(chunk_args)
                                buffer["args"] = "".join(parts)
                        elif chunk_args is not None:
                            buffer["args"] = chunk_args

                        buffer_name = buffer.get("name")
                        buffer_id = buffer.get("id")
                        if buffer_name is None:
                            continue

                        parsed_args = buffer.get("args")
                        if isinstance(parsed_args, str):
                            if not parsed_args:
                                continue
                            try:
                                parsed_args = json.loads(parsed_args)
                            except json.JSONDecodeError:
                                continue
                        elif parsed_args is None:
                            continue

                        if not isinstance(parsed_args, dict):
                            parsed_args = {"value": parsed_args}

                        # Flush text before tool call
                        if pending_text:
                            yield TextComplete(pending_text, current_message_id or "")
                            pending_text = ""
                            current_message_id = None

                        if (
                            buffer_id is not None
                            and buffer_id not in displayed_tool_ids
                        ):
                            displayed_tool_ids.add(buffer_id)
                            file_op_tracker.start_operation(
                                buffer_name, parsed_args, buffer_id
                            )
                            yield StatusUpdate(None)
                            yield ToolCallStart(buffer_id, buffer_name, parsed_args)
                            active_tool_calls[buffer_id] = (buffer_name, parsed_args)

                        tool_call_buffers.pop(buffer_key, None)

                if getattr(message, "chunk_position", None) == "last":
                    if pending_text:
                        yield TextComplete(pending_text, current_message_id or "")
                        pending_text = ""
                        current_message_id = None

            # Flush remaining text after stream
            if pending_text:
                yield TextComplete(pending_text, current_message_id or "")
                pending_text = ""
                current_message_id = None

            # ---- HITL handling ----
            if interrupt_occurred and pending_interrupts:
                for interrupt_id, hitl_request in pending_interrupts.items():
                    action_requests = hitl_request["action_requests"]

                    if auto_approve:
                        decisions: list[HITLDecision] = [
                            ApproveDecision(type="approve") for _ in action_requests
                        ]
                        hitl_response[interrupt_id] = {"decisions": decisions}
                        yield ToolsApproved()
                        _mark_file_ops_approved(file_op_tracker, action_requests)

                    elif resolve_hitl is not None:
                        decision = await resolve_hitl(action_requests, assistant_id)

                        if not isinstance(decision, dict):
                            decisions = [
                                RejectDecision(type="reject") for _ in action_requests
                            ]
                            hitl_response[interrupt_id] = {"decisions": decisions}
                            yield ToolsRejected()
                            return

                        decision_type = decision.get("type")

                        if decision_type == "auto_approve_all":
                            auto_approve = True
                            yield AutoApproveEnabled()
                            decisions = [
                                ApproveDecision(type="approve")
                                for _ in action_requests
                            ]
                            hitl_response[interrupt_id] = {"decisions": decisions}
                            yield ToolsApproved()
                            _mark_file_ops_approved(file_op_tracker, action_requests)

                        elif decision_type == "approve":
                            decisions = [
                                ApproveDecision(type="approve")
                                for _ in action_requests
                            ]
                            hitl_response[interrupt_id] = {"decisions": decisions}
                            yield ToolsApproved()
                            _mark_file_ops_approved(file_op_tracker, action_requests)

                        else:
                            decisions = [
                                RejectDecision(type="reject") for _ in action_requests
                            ]
                            hitl_response[interrupt_id] = {"decisions": decisions}
                            yield ToolsRejected()
                            return

                    else:
                        # No callback and not auto-approve — reject
                        decisions = [
                            RejectDecision(type="reject") for _ in action_requests
                        ]
                        hitl_response[interrupt_id] = {"decisions": decisions}
                        yield ToolsRejected()
                        return

                stream_input = Command(resume=hitl_response)
            else:
                break

        if captured_tokens:
            yield TokenUsage(captured_tokens)

    except (asyncio.CancelledError, KeyboardInterrupt):
        # Save interrupted state to graph (best-effort)
        try:
            interrupted_msg = _build_interrupted_ai_message(
                pending_text, active_tool_calls
            )
            if interrupted_msg:
                await agent.aupdate_state(config, {"messages": [interrupted_msg]})
            cancellation_msg = HumanMessage(
                content="[SYSTEM] Task interrupted by user. "
                "Previous operation was cancelled."
            )
            await agent.aupdate_state(config, {"messages": [cancellation_msg]})
        except Exception:
            logger.debug("Failed to save interrupted state", exc_info=True)
        raise
