import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, List

from agentic_base.evaluation import VectorStore
from agentic_base.evaluation.retrievers import Retriever
from agentic_base.evaluation.retrievers.dense import DenseRetriever
from fastmcp import FastMCP
from pydantic import BaseModel

from agentic_mcp.settings import EvidenceBuilderSettings, MCPRunnerSettings

_logger = logging.getLogger(__name__)


class MCPRunner(ABC):
    class RetrievedDocument(BaseModel):
        score: float
        page_content: str
        metadata: dict

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

    def _register_tools(self) -> None:
        @self.app.tool()
        async def retrieve(query: str) -> dict:
            _logger.info(f"Retrieve request | query='{query}' | vector_limit={self.config.retriever.limit}")
            response = await self.retriever.retrieve(
                query=query,
                limit=self.config.retriever.limit,
            )
            points = response.points
            _logger.info(f"Vector retrieval done | chunks={len(points)}")
            evidence = []
            for p in points:
                evidence.append(
                    {"score": p.score, "chunk": p.document.page_content}
                )  # retriever_limit * retriever_max_tokens
            return {
                "retrieval": {
                    "evidence": evidence,
                }
            }

    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:
    def __init__(self, config: EvidenceBuilderSettings):
        self.config = 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,
            }
        }
