"""Conversation management service."""

import contextlib
from collections.abc import AsyncIterator, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
from pydantic_ai.toolsets import AbstractToolset

from .message_store import MessageStore

if TYPE_CHECKING:
    from .usage_tracker import UsageTracker


@dataclass
class ConversationContext:
    """Context for a conversation."""

    conversation_id: str
    user_id: str
    messages: list[ModelMessage]
    metadata: dict[str, Any]


class ConversationService:
    """Service for managing multi-turn conversations.

    Handles the lifecycle of conversations including:
    - Creating new conversations
    - Loading conversation history
    - Running agents with history
    - Streaming responses
    - Persisting new messages

    Example:
        ```python
        service = ConversationService(message_store)

        # Start or continue a conversation
        async for chunk in service.chat(
            conversation_id="conv-123",
            user_id="user-456",
            prompt="Hello!",
            agent=router_agent,
            deps=deps,
        ):
            print(chunk, end="")
        ```
    """

    def __init__(
        self,
        message_store: MessageStore,
        usage_tracker: "UsageTracker | None" = None,
    ) -> None:
        """Initialize the conversation service.

        Args:
            message_store: Store for persisting messages.
            usage_tracker: Optional tracker for recording usage metrics.
        """
        self.message_store = message_store
        self.usage_tracker = usage_tracker

    async def get_context(
        self,
        conversation_id: str,
        user_id: str,
    ) -> ConversationContext:
        """Load conversation context.

        Args:
            conversation_id: Unique identifier for the conversation.
            user_id: The user's ID.

        Returns:
            ConversationContext with loaded messages.
        """
        messages = await self.message_store.load_messages(conversation_id)

        return ConversationContext(
            conversation_id=conversation_id,
            user_id=user_id,
            messages=messages,
            metadata={},
        )

    async def chat(
        self,
        conversation_id: str,
        user_id: str,
        prompt: str,
        agent: Agent[Any, Any],
        deps: Any,
        agent_name: str = "unknown",
        model: str = "unknown",
        toolsets: Sequence[AbstractToolset] | None = None,
    ) -> AsyncIterator[str]:
        """Run a chat turn with streaming response.

        Args:
            conversation_id: Unique identifier for the conversation.
            user_id: The user's ID.
            prompt: The user's message.
            agent: The agent to run.
            deps: Dependencies for the agent.
            agent_name: Name of the agent for usage tracking.
            model: Model name for usage tracking.
            toolsets: Optional additional toolsets (e.g., MCP servers).

        Yields:
            Chunks of the response text as they're generated.
        """
        # Load conversation history
        context = await self.get_context(conversation_id, user_id)

        # Run agent with streaming
        async with agent.run_stream(
            prompt,
            deps=deps,
            message_history=context.messages,
            toolsets=toolsets or [],
        ) as result:
            # Stream the response using stream_text() for proper deltas
            # (stream_output() returns accumulated output, not deltas)
            async for chunk in result.stream_text(delta=True):
                yield chunk

            # Save new messages after completion
            await self.message_store.save_messages(
                conversation_id=conversation_id,
                messages=result.new_messages(),
                user_id=user_id,
            )

            # Track usage if tracker is available (don't fail chat on tracking errors)
            if self.usage_tracker:
                with contextlib.suppress(Exception):
                    await self.usage_tracker.record_usage(
                        user_id=user_id,
                        agent_name=agent_name,
                        model=model,
                        usage=result.usage(),
                    )

    async def chat_sync(
        self,
        conversation_id: str,
        user_id: str,
        prompt: str,
        agent: Agent[Any, Any],
        deps: Any,
        agent_name: str = "unknown",
        model: str = "unknown",
        toolsets: Sequence[AbstractToolset] | None = None,
    ) -> str:
        """Run a chat turn and return the complete response.

        Args:
            conversation_id: Unique identifier for the conversation.
            user_id: The user's ID.
            prompt: The user's message.
            agent: The agent to run.
            deps: Dependencies for the agent.
            agent_name: Name of the agent for usage tracking.
            model: Model name for usage tracking.
            toolsets: Optional additional toolsets (e.g., MCP servers).

        Returns:
            The complete response text.
        """
        # Load conversation history
        context = await self.get_context(conversation_id, user_id)

        # Run agent
        result = await agent.run(
            prompt,
            deps=deps,
            message_history=context.messages,
            toolsets=toolsets or [],
        )

        # Save new messages
        await self.message_store.save_messages(
            conversation_id=conversation_id,
            messages=result.new_messages(),
            user_id=user_id,
        )

        # Track usage if tracker is available (don't fail chat on tracking errors)
        if self.usage_tracker:
            with contextlib.suppress(Exception):
                await self.usage_tracker.record_usage(
                    user_id=user_id,
                    agent_name=agent_name,
                    model=model,
                    usage=result.usage(),
                )

        # Return output as string
        output = result.output
        if hasattr(output, "model_dump_json"):
            return output.model_dump_json(indent=2)
        return str(output)

    async def clear_conversation(
        self,
        conversation_id: str,
    ) -> bool:
        """Clear all messages from a conversation.

        Args:
            conversation_id: Unique identifier for the conversation.

        Returns:
            True if conversation was cleared, False if not found.
        """
        return await self.message_store.delete_messages(conversation_id)

    async def list_conversations(
        self,
        user_id: str,
        limit: int = 20,
    ) -> list[dict[str, Any]]:
        """List recent conversations for a user.

        Args:
            user_id: The user's ID.
            limit: Maximum number of conversations to return.

        Returns:
            List of conversation metadata.
        """
        return await self.message_store.get_recent_conversations(user_id, limit)
