import logging
from typing import Any, Dict, Generic, List

from agentic_base.evaluation import VectorStore
from agentic_base.evaluation.retrievers import Retriever
from agentic_base.evaluation.retrievers.dense import DenseRetriever
from agentic_base.mixins import FromConfigMixin
from agentic_base.utils import load_class
from agentic_mcp.constants import TCEvidenceBuilder, TCMCPRunner
from fastmcp import FastMCP
from pydantic import BaseModel

_logger = logging.getLogger(__name__)


class MCPRunner(FromConfigMixin[TCMCPRunner], Generic[TCMCPRunner]):
    class RetrievedDocument(BaseModel):
        score: float
        page_content: str
        metadata: dict

    def __init__(self, config: TCMCPRunner) -> None:
        super().__init__(config)
        _logger.info("Initializing MCPRetrievalServer | name=mcp")
        self.app = FastMCP("mcp")
        self.retriever: Retriever[Any] = DenseRetriever.from_config(self.config.retriever.engine)
        self.evidence: EvidenceBuilder[Any] = load_class(self.config.evidence.module_path).from_config(
            self.config.evidence
        )
        self._register_tools()

    def _register_tools(self) -> None:
        @self.app.tool()
        async def retrieve(query: str) -> Dict[str, Any]:
            _logger.info(f"Retrieve request | query='{query}' | vector_limit={self.config.retriever.final_limit}")
            response = await self.retriever.retrieve(
                query=query,
                retrieve_limit=self.config.retriever.retrieve_limit,
                final_limit=self.config.retriever.final_limit,
            )
            points = response.points
            _logger.info(f"Vector retrieval done | chunks={len(points)}")
            return self.evidence.build(points)

    def run(self) -> None:
        _logger.info(
            f"Starting MCP server | entrypoint={self.config.module_path}",
        )
        self.app.run(
            transport=self.config.server.transport,
            host=self.config.server.host,
            port=self.config.server.port,
        )


class EvidenceBuilder(FromConfigMixin[TCEvidenceBuilder], Generic[TCEvidenceBuilder]):
    def __init__(self, config: TCEvidenceBuilder):
        super().__init__(config)

    def build(self, points: List[VectorStore.ScoredPoint]) -> Dict[str, Any]:
        _logger.info(f"Vector retrieval done | chunks={len(points)}")
        evidence = []
        for p in points:
            evidence.append(
                {"score": p.score, "page_content": p.document.page_content, "metadata": p.document.metadata}
            )
        return {
            "retrieval": {
                "evidence": evidence,
            }
        }
