"""Base stream handler for agent event processing.

Provides shared event processing logic for both TUI (Textual) and WebSocket
adapters. Platform-specific behavior is implemented by subclasses through
the on_* callback methods.
"""

from __future__ import annotations

import asyncio
import json
import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any

from langchain.agents.middleware.human_in_the_loop import (
    HITLRequest,
    HITLResponse,
)
from langchain_core.messages import HumanMessage, ToolMessage
from langgraph.types import Command, Interrupt
from pydantic import TypeAdapter, ValidationError

from ieum.file_ops import FileOpTracker
from ieum.approval_policy import decide_actions, normalize_approval_policy
from ieum.ui import format_tool_message_content

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


def normalize_token_usage(usage_data: dict | None) -> dict:
    """Normalize token usage from various LLM providers into a unified format.

    Supports:
    - OpenAI format: prompt_tokens / completion_tokens / total_tokens
    - Anthropic format: input_tokens / output_tokens
    - Ollama format: prompt_eval_count / eval_count
    - Generic format: any key matching *_tokens pattern

    Args:
        usage_data: Raw usage dict from an LLM provider response.

    Returns:
        Dict with ``input_tokens``, ``output_tokens``, and ``total_tokens``
        keys.  All values default to ``0`` when the source field is absent.
    """
    if not usage_data:
        return {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}

    input_tokens = (
        usage_data.get("input_tokens")  # Anthropic
        or usage_data.get("prompt_tokens")  # OpenAI / OpenRouter / Groq
        or usage_data.get("prompt_eval_count")  # Ollama
        or 0
    )
    output_tokens = (
        usage_data.get("output_tokens")  # Anthropic
        or usage_data.get("completion_tokens")  # OpenAI / OpenRouter / Groq
        or usage_data.get("eval_count")  # Ollama
        or 0
    )
    total = usage_data.get("total_tokens") or (input_tokens + output_tokens)

    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total,
    }

_HITL_REQUEST_ADAPTER = TypeAdapter(HITLRequest)

# Default timeout for pending futures (10 minutes)
APPROVAL_TIMEOUT_SECONDS = 600
INTERVIEW_TIMEOUT_SECONDS = 600


def _is_summarization_chunk(metadata: dict | None) -> bool:
    """Check if a message chunk is from summarization middleware.

    Args:
        metadata: The metadata dict from the stream chunk.

    Returns:
        Whether the chunk is from summarization and should be filtered.
    """
    if metadata is None:
        return False
    return metadata.get("lc_source") == "summarization"


class BaseStreamHandler(ABC):
    """Base class for handling agent stream events.

    Subclasses must implement the on_* callback methods to define
    platform-specific behavior (TUI widgets, WebSocket messages, etc.).

    The process_stream method contains the shared event processing logic
    that dispatches to the callbacks.
    """

    # --- Abstract callbacks that subclasses must implement ---

    @abstractmethod
    async def on_chat_token(self, token: str, ns_key: tuple) -> None:
        """Handle a streaming text token from the assistant.

        Args:
            token: The text token content.
            ns_key: Namespace key for tracking messages per subagent.
        """
        ...

    @abstractmethod
    async def on_tool_start(
        self, tool_id: str, tool_name: str, tool_input: dict, ns_key: tuple,
    ) -> None:
        """Handle start of a tool call.

        Args:
            tool_id: Unique tool call identifier.
            tool_name: Name of the tool being called.
            tool_input: Parsed arguments for the tool.
            ns_key: Namespace key.
        """
        ...

    @abstractmethod
    async def on_tool_end(
        self, tool_id: str, tool_name: str, output: str, status: str,
    ) -> None:
        """Handle completion of a tool call.

        Args:
            tool_id: Unique tool call identifier.
            tool_name: Name of the tool.
            output: Tool output string.
            status: Either 'success' or 'error'.
        """
        ...

    @abstractmethod
    async def on_file_changed(self, record: Any) -> None:
        """Handle a file operation result (diff display).

        Args:
            record: FileOpTracker record with diff information.
        """
        ...

    @abstractmethod
    async def on_thinking_token(self, token: str) -> None:
        """Handle a thinking/reasoning token from the model.

        Args:
            token: The thinking text content.
        """
        ...

    @abstractmethod
    async def on_thinking_end(self) -> None:
        """Handle end of thinking block."""
        ...

    @abstractmethod
    async def on_message_complete(self, ns_key: tuple, tokens: int) -> None:
        """Handle completion of an assistant message for a namespace.

        Args:
            ns_key: Namespace key.
            tokens: Token count for the message.
        """
        ...

    @abstractmethod
    async def on_human_message(self, ns_key: tuple) -> None:
        """Handle a human message in the stream (flush pending text).

        Args:
            ns_key: Namespace key.
        """
        ...

    @abstractmethod
    async def on_stream_error(self, error: str) -> None:
        """Handle a stream-level error.

        Args:
            error: Error description.
        """
        ...

    @abstractmethod
    async def on_interrupted(self) -> None:
        """Handle stream interruption by user."""
        ...

    @abstractmethod
    async def on_tool_rejected_cleanup(self) -> None:
        """Clean up tool state when execution is rejected or cancelled."""
        ...

    @abstractmethod
    async def on_rejection_message(self) -> None:
        """Display rejection message to user."""
        ...

    @abstractmethod
    async def request_approval(
        self, action_requests: list[dict], assistant_id: str | None,
    ) -> asyncio.Future:
        """Request approval from the user for tool calls.

        Args:
            action_requests: List of tool call action requests.
            assistant_id: Agent identifier.

        Returns:
            Future that resolves with the approval decision.
        """
        ...

    @abstractmethod
    async def request_interview(
        self, interview_id: str, questions: list[dict],
    ) -> asyncio.Future:
        """Request interview answers from the user.

        Args:
            interview_id: Unique interview identifier.
            questions: List of question dicts to present.

        Returns:
            Future that resolves with the interview answers.
        """
        ...

    @abstractmethod
    async def on_approval_auto_approved(self) -> None:
        """Handle auto-approve being enabled for the session."""
        ...

    @abstractmethod
    async def on_pre_stream(self) -> None:
        """Called before streaming begins (e.g., show spinner)."""
        ...

    @abstractmethod
    async def on_post_stream(self, tokens: int) -> None:
        """Called after streaming completes (e.g., update token display).

        Args:
            tokens: Total captured token count.
        """
        ...

    @abstractmethod
    async def on_interview_complete(self, interview_id: str) -> None:
        """Handle interview completion notification.

        Args:
            interview_id: The interview that was completed.
        """
        ...

    # --- Hooks for optional override ---

    async def on_token_usage(self, total_tokens: int) -> None:
        """Handle token usage information from a message.

        Args:
            total_tokens: Total token count extracted from the message.
        """
        pass

    async def on_tool_call_flush_text(self, ns_key: tuple) -> None:
        """Called before a tool call to flush any pending text.

        Args:
            ns_key: Namespace key.
        """
        pass

    async def on_tool_running(self) -> None:
        """Called when tool starts running after approval (show animation)."""
        pass

    async def on_values_update(self, data: dict) -> None:
        """Handle a values stream update (full state snapshot).

        Args:
            data: State data dict.
        """
        pass

    async def on_debug_event(self, data: dict) -> None:
        """Handle a debug stream event.

        Args:
            data: Debug event data.
        """
        pass

    def should_process_values_stream(self) -> bool:
        """Whether to process values stream mode. Override in WS adapter."""
        return False

    def should_process_debug_stream(self) -> bool:
        """Whether to process debug stream mode. Override in WS adapter."""
        return False

    # --- Shared stream processing logic ---

    async def process_stream(
        self,
        agent: Any,
        stream_input: dict | Command,
        config: dict,
        file_op_tracker: FileOpTracker,
        session_auto_approve: bool,
        assistant_id: str | None,
        approval_policy: str = "manual",
        on_auto_approve_enabled: Any | None = None,
        on_mark_hitl_approved: Any | None = None,
    ) -> tuple[bool, dict | Command | None, int]:
        """Process a single iteration of the agent stream.

        This handles one pass through agent.astream(), processing all events,
        then handles HITL interrupts. Callers should loop calling this method
        until should_continue is False.

        Args:
            agent: The LangGraph agent.
            stream_input: Input for the stream (messages or Command(resume=...)).
            config: LangGraph config dict.
            file_op_tracker: Tracker for file operations.
            session_auto_approve: Whether auto-approve is enabled.
            assistant_id: Agent identifier.
            approval_policy: Approval preset policy.
            on_auto_approve_enabled: Callback when auto-approve becomes enabled.
            on_mark_hitl_approved: Callback to mark file ops as HITL approved.

        Returns:
            Tuple of (should_continue, next_stream_input, captured_tokens).
            If should_continue is True, call again with next_stream_input.
            If should_continue is False, the stream is complete.
        """
        interrupt_occurred = False
        interview_occurred = False
        hitl_response: dict[str, HITLResponse] = {}
        suppress_resumed_output = False
        pending_interrupts: dict[str, HITLRequest] = {}
        pending_interviews: dict[str, dict] = {}

        displayed_tool_ids: set[str] = set()
        tool_call_buffers: dict[str | int, dict] = {}
        pending_text_by_namespace: dict[tuple, str] = {}
        captured_tokens = 0

        stream_modes = ["messages", "updates"]
        if self.should_process_values_stream():
            stream_modes.append("values")
        if self.should_process_debug_stream():
            stream_modes.append("debug")

        async for chunk in agent.astream(
            stream_input,
            stream_mode=stream_modes,
            subgraphs=True,
            config=config,
            durability="exit",
        ):
            if not isinstance(chunk, tuple) or len(chunk) != 3:
                continue

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

            # Handle VALUES stream (full state snapshots)
            if current_stream_mode == "values":
                if is_main_agent and isinstance(data, dict):
                    await self.on_values_update(data)
                continue

            # Handle DEBUG stream
            if current_stream_mode == "debug":
                if is_main_agent and isinstance(data, dict):
                    await self.on_debug_event(data)
                continue

            # Handle UPDATES stream - for interrupts and todos
            if current_stream_mode == "updates":
                if not isinstance(data, dict):
                    continue

                if "__interrupt__" in data:
                    interrupts: list[Interrupt] = data["__interrupt__"]
                    if interrupts:
                        for interrupt_obj in interrupts:
                            # Check for interview interrupts (ask_user_question)
                            if (
                                isinstance(interrupt_obj.value, dict)
                                and interrupt_obj.value.get("type") == "ask_user_question"
                            ):
                                pending_interviews[interrupt_obj.id] = interrupt_obj.value
                                interview_occurred = True
                                continue

                            try:
                                validated_request = _HITL_REQUEST_ADAPTER.validate_python(
                                    interrupt_obj.value
                                )
                                pending_interrupts[interrupt_obj.id] = validated_request
                                interrupt_occurred = True
                            except ValidationError:
                                raise

            # Handle MESSAGES stream - for content and tool calls
            elif current_stream_mode == "messages":
                # Skip subagent outputs
                if not is_main_agent:
                    continue

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

                message, _metadata = data

                # Filter out summarization LLM output
                if _is_summarization_chunk(_metadata):
                    continue

                if isinstance(message, HumanMessage):
                    pending_text = pending_text_by_namespace.get(ns_key, "")
                    if pending_text:
                        await self.on_human_message(ns_key)
                        pending_text_by_namespace[ns_key] = ""
                    continue

                if isinstance(message, ToolMessage):
                    tool_name = getattr(message, "name", "")
                    tool_status = getattr(message, "status", "success")
                    tool_content = format_tool_message_content(message.content)
                    record = file_op_tracker.complete_with_message(message)

                    tool_id = getattr(message, "tool_call_id", None)
                    if tool_id:
                        output_str = str(tool_content) if tool_content else ""
                        await self.on_tool_end(tool_id, tool_name, output_str, tool_status)

                    # Send file change notification
                    if record and record.diff:
                        await self.on_file_changed(record)
                    continue

                # Extract token usage
                captured_tokens = self._extract_tokens(message, captured_tokens)

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

                # Process content blocks
                for block in message.content_blocks:
                    block_type = block.get("type")

                    # Handle thinking blocks
                    if block_type == "thinking":
                        thinking_text = block.get("thinking", "")
                        if thinking_text:
                            await self.on_thinking_token(thinking_text)
                        continue

                    if block_type == "text":
                        text = block.get("text", "")
                        if text:
                            pending_text = pending_text_by_namespace.get(ns_key, "")
                            pending_text += text
                            pending_text_by_namespace[ns_key] = pending_text
                            await self.on_chat_token(text, ns_key)

                    elif block_type in ("tool_call_chunk", "tool_call"):
                        self._process_tool_call_block(
                            block, tool_call_buffers,
                        )

                        # Check if we have a complete tool call to display
                        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.get(buffer_key)
                        if buffer is None:
                            continue

                        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 pending text before tool call
                        pending_text = pending_text_by_namespace.get(ns_key, "")
                        if pending_text:
                            await self.on_tool_call_flush_text(ns_key)
                            pending_text_by_namespace[ns_key] = ""

                        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)
                            await self.on_tool_start(buffer_id, buffer_name, parsed_args, ns_key)

                        tool_call_buffers.pop(buffer_key, None)

                if getattr(message, "chunk_position", None) == "last":
                    await self.on_thinking_end()
                    pending_text = pending_text_by_namespace.get(ns_key, "")
                    if pending_text:
                        await self.on_message_complete(ns_key, captured_tokens)
                        pending_text_by_namespace[ns_key] = ""

        # Flush remaining text from all namespaces
        await self.on_thinking_end()
        for ns_key, pending_text in list(pending_text_by_namespace.items()):
            if pending_text:
                await self.on_message_complete(ns_key, captured_tokens)
        pending_text_by_namespace.clear()

        # Handle HITL after stream completes
        if interrupt_occurred:
            any_rejected = False
            normalized_policy = normalize_approval_policy(approval_policy)

            for interrupt_id, hitl_request in pending_interrupts.items():
                action_requests = hitl_request["action_requests"]

                if session_auto_approve:
                    decisions = [{"type": "approve"} for _ in action_requests]
                    hitl_response[interrupt_id] = {"decisions": decisions}
                    await self.on_tool_running()
                else:
                    policy_decisions = decide_actions(normalized_policy, action_requests)
                    if policy_decisions is not None:
                        decisions = [{"type": decision.decision} for decision in policy_decisions]
                        hitl_response[interrupt_id] = {"decisions": decisions}
                        approved_indexes = [
                            idx
                            for idx, decision in enumerate(decisions)
                            if decision.get("type") == "approve"
                        ]
                        if approved_indexes:
                            await self.on_tool_running()
                            if on_mark_hitl_approved:
                                for idx in approved_indexes:
                                    on_mark_hitl_approved(action_requests[idx])
                        if len(approved_indexes) != len(decisions):
                            await self.on_tool_rejected_cleanup()
                            any_rejected = True
                            break
                        continue

                    future = await self.request_approval(action_requests, assistant_id)
                    try:
                        decision = await asyncio.wait_for(
                            future, timeout=APPROVAL_TIMEOUT_SECONDS,
                        )
                    except TimeoutError:
                        logger.warning(
                            "Approval timed out after %d seconds, auto-rejecting",
                            APPROVAL_TIMEOUT_SECONDS,
                        )
                        decision = {"type": "reject"}

                    if isinstance(decision, dict):
                        decision_type = decision.get("type")

                        if decision_type == "auto_approve_all":
                            if on_auto_approve_enabled:
                                on_auto_approve_enabled()
                            decisions = [{"type": "approve"} for _ in action_requests]
                            await self.on_tool_running()
                            await self.on_approval_auto_approved()
                            if on_mark_hitl_approved:
                                for action_request in action_requests:
                                    on_mark_hitl_approved(action_request)

                        elif decision_type == "approve":
                            decisions = [{"type": "approve"} for _ in action_requests]
                            await self.on_tool_running()
                            if on_mark_hitl_approved:
                                for action_request in action_requests:
                                    on_mark_hitl_approved(action_request)

                        elif decision_type == "reject":
                            decisions = [{"type": "reject"} for _ in action_requests]
                            await self.on_tool_rejected_cleanup()
                            any_rejected = True
                        else:
                            decisions = [{"type": "reject"} for _ in action_requests]
                            any_rejected = True
                    else:
                        decisions = [{"type": "reject"} for _ in action_requests]
                        any_rejected = True

                    hitl_response[interrupt_id] = {"decisions": decisions}

                    if any_rejected:
                        break

            suppress_resumed_output = any_rejected

        # Handle interview interrupts
        if interview_occurred and pending_interviews:
            for iview_id, iview_data in pending_interviews.items():
                questions = iview_data.get("questions", [])
                future = await self.request_interview(iview_id, questions)
                try:
                    answers = await asyncio.wait_for(
                        future, timeout=INTERVIEW_TIMEOUT_SECONDS,
                    )
                except TimeoutError:
                    logger.warning(
                        "Interview timed out after %d seconds, auto-skipping",
                        INTERVIEW_TIMEOUT_SECONDS,
                    )
                    answers = {}

                if isinstance(answers, dict) and answers.get("_cancelled"):
                    hitl_response[iview_id] = {"_cancelled": True}
                    suppress_resumed_output = True
                else:
                    hitl_response[iview_id] = answers

                await self.on_interview_complete(iview_id)

            interview_occurred = False
            pending_interviews.clear()

        # Determine whether to continue streaming
        if (interrupt_occurred or hitl_response) and hitl_response:
            if suppress_resumed_output:
                await self.on_rejection_message()
                return False, None, captured_tokens

            next_input = Command(resume=hitl_response)
            return True, next_input, captured_tokens
        else:
            return False, None, captured_tokens

    def _extract_tokens(self, message: Any, current_max: int) -> int:
        """Extract token usage from a message, returning the new maximum.

        Uses :func:`normalize_token_usage` to handle multiple LLM provider
        response formats (OpenAI, Anthropic, Ollama, Groq, OpenRouter).

        Also stores the input/output breakdown in ``_token_breakdown`` for
        cost calculation by downstream consumers.

        Args:
            message: The LLM message to extract tokens from.
            current_max: Current maximum token count.

        Returns:
            Updated maximum token count.
        """
        # Try usage_metadata first (LangChain standard attribute)
        if hasattr(message, "usage_metadata") and message.usage_metadata:
            normalized = normalize_token_usage(message.usage_metadata)
            total = normalized["total_tokens"]
            if total:
                self._token_breakdown = normalized
                return max(current_max, total)

        # Fall back to response_metadata (provider-specific data)
        if hasattr(message, "response_metadata") and message.response_metadata:
            resp = message.response_metadata
            # First try top-level keys (Ollama, OpenRouter)
            normalized = normalize_token_usage(resp)
            if normalized["total_tokens"]:
                self._token_breakdown = normalized
                return max(current_max, normalized["total_tokens"])
            # Then try nested "usage" dict (Groq)
            if "usage" in resp and isinstance(resp["usage"], dict):
                normalized = normalize_token_usage(resp["usage"])
                if normalized["total_tokens"]:
                    self._token_breakdown = normalized
                    return max(current_max, normalized["total_tokens"])

        return current_max

    @staticmethod
    def _process_tool_call_block(
        block: dict, tool_call_buffers: dict[str | int, dict],
    ) -> None:
        """Process a tool_call_chunk or tool_call content block into the buffer.

        Accumulates streamed tool call arguments across multiple chunks.

        Args:
            block: The content block from the message.
            tool_call_buffers: Mutable dict of accumulated tool call data.
        """
        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
