import asyncio
from collections.abc import AsyncGenerator

from rich.console import Console

from alloy_runtime_types.dtos.generation import GenerateTextStreamChunk


class ExecutionState:
    def __init__(self, index: int, agent_identifier: str | None = None):
        self.index = index
        self.agent_identifier = agent_identifier
        self.thinking_text = ""
        self.content_text = ""
        self.execution_id: str | None = None
        self.session_id: str | None = None
        self.is_complete = False
        self.has_error = False
        self.error_message: str | None = None

    def process_chunk(self, chunk: GenerateTextStreamChunk) -> None:
        if self.execution_id is None:
            self.execution_id = str(chunk.execution_id)

        if chunk.chunk_type == "thinking":
            if chunk.reasoning_content:
                self.thinking_text += chunk.reasoning_content
        elif chunk.chunk_type == "metadata":
            if chunk.metadata and "chat_session_id" in chunk.metadata:
                self.session_id = chunk.metadata["chat_session_id"]
        elif chunk.content:
            self.content_text += chunk.content


class ConcurrentRenderer:
    def __init__(
        self,
        console: Console,
        execution_count: int,
        agent_identifiers: list[str] | None = None,
    ):
        self.console = console
        self.execution_count = min(max(execution_count, 1), 10)
        self._states = [
            ExecutionState(
                i,
                agent_identifier=(
                    agent_identifiers[i]
                    if agent_identifiers and i < len(agent_identifiers)
                    else None
                ),
            )
            for i in range(self.execution_count)
        ]

    async def render_streams(
        self,
        streams: list[AsyncGenerator[GenerateTextStreamChunk, None]],
        show_thinking: bool = True,
    ) -> list[tuple[str, str, str | None, str | None]]:
        if len(streams) != self.execution_count:
            raise ValueError(
                f"Expected {self.execution_count} streams, got {len(streams)}"
            )

        await asyncio.gather(
            *(
                self._consume_stream(i, stream, show_thinking)
                for i, stream in enumerate(streams)
            ),
            return_exceptions=False,
        )

        return [
            (
                state.content_text,
                state.thinking_text,
                state.error_message,
                state.session_id,
            )
            for state in self._states
        ]

    async def _consume_stream(
        self,
        index: int,
        stream: AsyncGenerator[GenerateTextStreamChunk, None],
        show_thinking: bool,
    ) -> None:
        state = self._states[index]
        try:
            async for chunk in stream:
                if show_thinking or chunk.chunk_type != "thinking":
                    state.process_chunk(chunk)
            state.is_complete = True
        except Exception as error:
            state.has_error = True
            state.error_message = str(error)
            state.is_complete = True

    def _tail_text(self, text: str, max_visual_lines: int, width: int) -> str:
        _ = max_visual_lines
        _ = width
        return text
