from typing import Any

from agent.agent import Agent
from agent.agent_client import AgentClient, ChatOutcome, SkillsAgentClient
from agent.champ_client import ChampAgentClient
from agent.fake_client import FakeAgentClient
from agent.frontier_client import (
    GEMINI_CONSERVATIVE_TEMPERATURE,
    GEMINI_CREATIVE_TEMPERATURE,
    GeminiAgentClient,
    GptAgentClient,
)
from classes.base_models import ChatRequest
from providers.protocol import ChatProvider

SessionId = str
ConversationId = str
ModelType = str


class SessionRouter:
    """Per-(session, conversation) registry of AgentClient instances."""

    def __init__(
        self,
        wiki_skills_agent: Agent,
        wiki_skills_agent_short: Agent,
        champ_provider: ChatProvider,
        champ_vector_store: Any,
        openai_provider: ChatProvider,
        openai_model_id: str,
        gemini_provider: ChatProvider,
        gemini_model_id: str,
        fake_provider: ChatProvider,
    ) -> None:
        self.session_clients: dict[tuple[SessionId, ConversationId], AgentClient] = {}
        self.wiki_skills_agent = wiki_skills_agent
        self.wiki_skills_agent_short = wiki_skills_agent_short
        self.champ_provider = champ_provider
        self.champ_vector_store = champ_vector_store
        self.openai_provider = openai_provider
        self.openai_model_id = openai_model_id
        self.gemini_provider = gemini_provider
        self.gemini_model_id = gemini_model_id
        self.fake_provider = fake_provider

    def delete_session(self, session_id: SessionId) -> None:
        """Drop every AgentClient (and its ConversationHistory) for this session."""
        keys = [k for k in self.session_clients if k[0] == session_id]
        for k in keys:
            del self.session_clients[k]

    def send(
        self, payload: ChatRequest, documents: dict[str, str] | None = None
    ) -> ChatOutcome:
        # TODO: session/conversation ids should be created backend-side, not client-side.
        key = (payload.session_id, payload.conversation_id)

        agent_client = self.session_clients.get(key)
        if agent_client is None:
            agent_client = self._create_agent_client(payload.model_type)
            self.session_clients[key] = agent_client

        return agent_client.call(
            payload.human_message, lang=payload.lang, documents=documents
        )

    def _create_agent_client(self, model_type: str) -> AgentClient:
        if model_type == "skills_wiki":
            return SkillsAgentClient(agent=self.wiki_skills_agent)
        if model_type == "skills_wiki_short":
            return SkillsAgentClient(agent=self.wiki_skills_agent_short)
        if model_type == "champ":
            return ChampAgentClient(
                provider=self.champ_provider,
                vector_store=self.champ_vector_store,
            )
        if model_type == "openai":
            return GptAgentClient(
                provider=self.openai_provider,
                model_id=self.openai_model_id,
            )
        if model_type == "google-conservative":
            return GeminiAgentClient(
                provider=self.gemini_provider,
                model_id=self.gemini_model_id,
                temperature=GEMINI_CONSERVATIVE_TEMPERATURE,
            )
        if model_type == "google-creative":
            return GeminiAgentClient(
                provider=self.gemini_provider,
                model_id=self.gemini_model_id,
                temperature=GEMINI_CREATIVE_TEMPERATURE,
            )
        if model_type == "fake":
            return FakeAgentClient(provider=self.fake_provider)
        raise RuntimeError(f"Unknown model_type: {model_type}")