"""Streaming response handling for Kollabor CLI LLM service.

Handles streaming chunk processing, thinking content display,
and LLM API call orchestration. Extracted from LLMService as part of
the llm_service.py decomposition (Phase B).

Phase 8A: Thinking parsing and display delegated to kollabor_ai and kollabor_tui.
"""

import asyncio
import logging
from typing import Callable, List, Optional

from kollabor_tui.status.core_widgets import get_token_io_state

# Delegated thinking parsing (kollabor_ai)
from kollabor_ai.streaming_thinking_parser import (
    StreamingThinkingParser,
    StreamingThinkingState,
)

# Delegated thinking display formatting (kollabor_tui)
from kollabor_tui.thinking_display import ThinkingDisplayFormatter

logger = logging.getLogger(__name__)


class StreamingHandler:
    """Handles streaming LLM responses and thinking content display.

    Responsibilities:
    - Process streaming chunks from LLM API
    - Parse and display thinking content (delegated to kollabor_ai + kollabor_tui)
    - Stream response content to message renderer
    - Orchestrate LLM API calls via APICommunicationService
    - Clean up streaming state after completion/error
    """

    def __init__(self, api_service, message_display_service, renderer):
        """Initialize the streaming handler.

        Args:
            api_service: APICommunicationService for making LLM API calls
            message_display_service: MessageDisplayService for streaming output
            renderer: TerminalRenderer for thinking display updates
        """
        self.api_service = api_service
        self.message_display_service = message_display_service
        self.renderer = renderer

        # Streaming state (delegates to kollabor_ai + kollabor_tui)
        self._thinking_parser = StreamingThinkingParser()
        self._thinking_formatter = ThinkingDisplayFormatter()
        self._response_started = False

    async def call_llm(
        self,
        conversation_history,
        max_history: int,
        native_tools: Optional[List[dict]],
        mcp_discovery_complete: asyncio.Event,
        is_cancelled_fn: Callable[[], bool],
    ) -> str:
        """Make API call to LLM using APICommunicationService.

        Args:
            conversation_history: Current conversation history
            max_history: Maximum history messages to send
            native_tools: Native tool definitions for function calling (or None)
            mcp_discovery_complete: Event signaling MCP discovery is done
            is_cancelled_fn: Callable that returns True if request is cancelled

        Returns:
            LLM response string

        Raises:
            asyncio.CancelledError: If request was cancelled by user
        """
        # Reset streaming state for new request (delegates to packages)
        self._thinking_parser.reset()
        self._thinking_formatter.reset()
        self._response_started = False

        # Check for cancellation before starting
        if is_cancelled_fn():
            logger.info("API call cancelled before starting")
            raise asyncio.CancelledError("Request cancelled by user")

        # Wait for MCP discovery to complete (prevent race condition on first call)
        try:
            await asyncio.wait_for(mcp_discovery_complete.wait(), timeout=5.0)
        except asyncio.TimeoutError:
            logger.warning(
                "MCP discovery did not complete within 5 seconds - "
                "proceeding with available tools (MCP tools may not be available)"
            )

        # Delegate to API communication service
        try:
            return await self.api_service.call_llm(
                conversation_history=conversation_history,
                max_history=max_history,
                streaming_callback=self.handle_chunk,
                tools=native_tools,
            )
        except asyncio.CancelledError:
            logger.info("LLM API call was cancelled")
            self.cleanup()
            raise
        except Exception as e:
            logger.error(f"LLM API call failed: {e}")
            self.cleanup()
            raise

    async def handle_chunk(self, chunk: str) -> None:
        """Handle streaming content chunk from API.

        Args:
            chunk: Content chunk from streaming API response
        """
        # Delegate to kollabor_ai StreamingThinkingParser
        result = self._thinking_parser.parse_chunk(chunk)

        # Stream response content chunks
        for content in result.response_content:
            self._stream_response_chunk(content)

        # Display thinking content (delegated to kollabor_tui)
        if result.thinking_content:
            display_text = self._thinking_formatter.format(
                result.thinking_content, final=result.thinking_complete
            )
            if display_text:
                self.renderer.update_thinking(True, display_text)

        # Handle thinking complete state (switch to receiving mode)
        if result.thinking_complete:
            self.renderer.update_thinking(True, "Receiving...")

    def cleanup(self) -> None:
        """Clean up streaming state after request completion or failure.

        This ensures streaming state is properly reset even if errors occur.
        """
        self._thinking_parser.reset()
        self._thinking_formatter.reset()
        self._response_started = False

        # End streaming session in message display service if active
        if self.message_display_service.is_streaming_active():
            self.message_display_service.end_streaming_response()

        logger.debug("Cleaned up streaming state")

    def _stream_response_chunk(self, chunk: str) -> None:
        """Stream a response chunk in real-time to the message renderer.

        Args:
            chunk: Response content chunk to stream immediately
        """
        # Handle empty chunks gracefully
        if not chunk or not chunk.strip():
            return

        # Initialize streaming response if this is the first chunk
        if not self._response_started:
            self.message_display_service.message_coordinator.start_streaming_response()
            self._response_started = True
            # Start receiving mode for token I/O widget
            get_token_io_state().start_receiving()
            self.renderer.update_thinking(True, "Receiving...")

        # Count tokens as they stream in (real-time, not fake animation)
        get_token_io_state().add_chunk(len(chunk))

        # Stream the chunk through the message coordinator (proper architecture)
        self.message_display_service.message_coordinator.write_streaming_chunk(chunk)
