"""Streaming text renderer with thinking block support."""

import sys
from collections.abc import AsyncGenerator

from pydantic import BaseModel
from rich.console import Console

from alloy_runtime_types.dtos.generation import GenerateTextStreamChunk
from cli.infrastructure.output import OutputService
from cli.infrastructure.renderers.entity_renderer import FieldConfig


class StreamRenderResult(BaseModel):
    content: str
    thinking: str | None = None
    session_id: str | None = None


class StreamRenderer:
    """Renders streaming text generation with thinking block support."""

    def __init__(self, console: Console, stderr_console: Console | None = None):
        """Initialize stream renderer.

        Args:
            console: Rich console for output
        """
        self.console = console
        self.stderr_console = stderr_console or Console(file=sys.stderr)
        self._thinking_text = ""
        self._content_text = ""
        self._session_id: str | None = None
        self._in_thinking_block = False

    async def render_stream(
        self,
        stream: AsyncGenerator[GenerateTextStreamChunk, None],
        show_thinking: bool = True,
    ) -> tuple[str, str, str | None]:
        """Render a streaming generation.

        Args:
            stream: Async generator of stream chunks
            show_thinking: Whether to display thinking blocks (default True)

        Returns:
            Tuple of (content_text, thinking_text, session_id) - accumulated outputs
        """
        self._thinking_text = ""
        self._content_text = ""
        self._session_id = None
        self._in_thinking_block = False

        async for chunk in stream:
            self._process_chunk(chunk, show_thinking)

        # Close thinking block if still open at end of stream
        if self._in_thinking_block:
            self._close_thinking_block()

        return self._content_text, self._thinking_text, self._session_id

    def present_final_output(
        self, content: str, thinking: str, session_id: str | None, show_thinking: bool
    ) -> None:
        result = StreamRenderResult(
            content=content,
            thinking=thinking or None if show_thinking else None,
            session_id=session_id,
        )
        fields = [FieldConfig("content", "Output")]
        if show_thinking and thinking:
            fields.append(FieldConfig("thinking", "Thinking"))
        if session_id:
            fields.append(FieldConfig("session_id", "Session ID"))
        OutputService.get().entity(result, "Generated Text", fields, raw_payload=result)

    def _process_chunk(
        self, chunk: GenerateTextStreamChunk, show_thinking: bool
    ) -> None:
        """Process a stream chunk.

        Args:
            chunk: Stream chunk to process
            show_thinking: Whether to display thinking text
        """
        if chunk.chunk_type == "thinking":
            if chunk.reasoning_content:
                self._thinking_text += chunk.reasoning_content
                if show_thinking:
                    if not self._in_thinking_block:
                        self._open_thinking_block()
                    self.stderr_console.out(
                        chunk.reasoning_content, style="red", end=""
                    )

        elif chunk.chunk_type == "metadata":
            if chunk.metadata and "chat_session_id" in chunk.metadata:
                self._session_id = chunk.metadata["chat_session_id"]

        else:
            if self._in_thinking_block:
                self._close_thinking_block()

            if chunk.content:
                self._content_text += chunk.content

    def _open_thinking_block(self) -> None:
        """Print opening thinking tag to stderr."""
        self.stderr_console.out("\n<thinking>\n", style="red", end="")
        self._in_thinking_block = True

    def _close_thinking_block(self) -> None:
        """Print closing thinking tag to stderr."""
        if not self._thinking_text.endswith("\n"):
            self.stderr_console.out("\n", end="")
        self.stderr_console.out("</thinking>\n\n", style="red", end="")
        self._in_thinking_block = False
