"""bio_mcp_base.py — clase base común de todos los servidores MCP de biología."""
import asyncio
import json
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import wraps
from typing import Any, Callable

import requests
import structlog
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from prometheus_client import Counter, Histogram, start_http_server


log = structlog.get_logger()


@dataclass
class RateLimitConfig:
    """Política de límite de solicitudes."""
    max_requests_per_second: float = 3.0
    max_requests_per_hour: int | None = None
    burst: int = 5


class RateLimiter:
    """Token bucket rate limiter (async)."""

    def __init__(self, config: RateLimitConfig):
        self.config = config
        self.tokens = float(config.burst)
        self.last_refill = time.time()
        self.lock = asyncio.Lock()

    async def acquire(self) -> None:
        async with self.lock:
            now = time.time()
            elapsed = now - self.last_refill
            self.tokens = min(
                self.config.burst,
                self.tokens + elapsed * self.config.max_requests_per_second,
            )
            self.last_refill = now
            if self.tokens < 1.0:
                wait_time = (1.0 - self.tokens) / self.config.max_requests_per_second
                await asyncio.sleep(wait_time)
                self.tokens = 0.0
            else:
                self.tokens -= 1.0


# Métricas de Prometheus (compartidas por todos los servidores)
TOOL_CALLS = Counter(
    "mcp_tool_calls_total",
    "Número de llamadas a herramientas MCP",
    ["server", "tool", "status"],
)
TOOL_LATENCY = Histogram(
    "mcp_tool_latency_seconds",
    "Latencia de herramientas MCP (segundos)",
    ["server", "tool"],
    buckets=[0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0],
)


def instrument(server_name: str, tool_name: str):
    """Decorador que registra la latencia y la tasa de éxito de las llamadas a herramientas."""
    def decorator(func: Callable):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            logger = log.bind(server=server_name, tool=tool_name)
            start = time.perf_counter()
            try:
                result = await func(*args, **kwargs)
                elapsed = time.perf_counter() - start
                TOOL_CALLS.labels(server=server_name, tool=tool_name, status="success").inc()
                TOOL_LATENCY.labels(server=server_name, tool=tool_name).observe(elapsed)
                logger.info("tool_success", latency_ms=elapsed * 1000)
                return result
            except Exception as e:
                elapsed = time.perf_counter() - start
                TOOL_CALLS.labels(server=server_name, tool=tool_name, status="failure").inc()
                TOOL_LATENCY.labels(server=server_name, tool=tool_name).observe(elapsed)
                logger.error("tool_failure", latency_ms=elapsed * 1000, error=str(e))
                raise
        return wrapper
    return decorator


class BioMCPServerBase(ABC):
    """Clase base de todos los servidores MCP de biología."""

    def __init__(
        self,
        server_name: str,
        rate_limit_config: RateLimitConfig,
        max_retries: int = 3,
        prometheus_port: int | None = 9090,
    ):
        self.server = Server(server_name)
        self.server_name = server_name
        self.rate_limiter = RateLimiter(rate_limit_config)
        self.max_retries = max_retries
        if prometheus_port:
            try:
                start_http_server(prometheus_port)
            except OSError:
                pass  # Omite si el puerto ya está en uso
        self._register_handlers()

    def _register_handlers(self):
        """Registra los handlers estándar de MCP."""
        @self.server.list_tools()
        async def _list() -> list[Tool]:
            return self.tools

        @self.server.call_tool()
        async def _call(name: str, arguments: dict) -> list[TextContent]:
            try:
                result = await self.dispatch_tool(name, arguments)
                return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False, indent=2))]
            except requests.RequestException as e:
                log.error("api_request_failed", tool=name, error=str(e))
                return [TextContent(type="text", text=json.dumps({"error": f"Falló la solicitud a la API: {e}"}))]
            except Exception as e:
                log.error("tool_execution_error", tool=name, error=str(e), exc_info=True)
                return [TextContent(type="text", text=json.dumps({"error": str(e)}))]

    @property
    @abstractmethod
    def tools(self) -> list[Tool]:
        """Lista de herramientas que expone este servidor."""
        ...

    @abstractmethod
    async def dispatch_tool(self, name: str, arguments: dict) -> Any:
        """Despacha por nombre de herramienta."""
        ...

    async def http_get_with_retry(
        self,
        url: str,
        params: dict | None = None,
        timeout: int = 30,
    ) -> requests.Response:
        """Solicitud GET con límite de solicitudes y reintentos."""
        for attempt in range(self.max_retries):
            await self.rate_limiter.acquire()
            try:
                resp = requests.get(url, params=params, timeout=timeout)
                if resp.status_code == 429:  # Rate limit hit
                    wait_time = 2 ** attempt
                    log.warning("rate_limited", wait=wait_time, attempt=attempt, url=url)
                    await asyncio.sleep(wait_time)
                    continue
                resp.raise_for_status()
                return resp
            except requests.Timeout:
                if attempt == self.max_retries - 1:
                    raise
                await asyncio.sleep(2 ** attempt)
        raise RuntimeError(f"Se superó el máximo de reintentos: {url}")

    async def run_stdio(self):
        """stdio transport."""
        async with stdio_server() as (read, write):
            await self.server.run(read, write, self.server.create_initialization_options())

"""pubmed_mcp.py — servidor MCP de PubMed E-utilities preparado para producción."""
import asyncio
import os
from xml.etree import ElementTree as ET
from typing import Any

from mcp.types import Tool
from bio_mcp_base import BioMCPServerBase, RateLimitConfig, instrument


EUTILS_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"


class PubMedMCPServer(BioMCPServerBase):
    def __init__(self, api_key: str | None = None):
        # Con una clave de API de NCBI admite 10 solicitudes/s; sin ella, 3 solicitudes/s.
        rate = 10.0 if api_key else 3.0
        super().__init__(
            server_name="pubmed-mcp",
            rate_limit_config=RateLimitConfig(max_requests_per_second=rate, burst=int(rate)),
        )
        self.api_key = api_key

    @property
    def tools(self) -> list[Tool]:
        return [
            Tool(
                name="search_pubmed",
                description="Devuelve una lista de PMID a partir de una consulta de PubMed; permite indicar intervalo de fechas y orden.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {"type": "string"},
                        "max_results": {"type": "integer", "default": 20},
                        "date_range_years": {"type": "integer", "default": 5},
                        "sort": {"type": "string", "enum": ["relevance", "pub_date"], "default": "relevance"},
                    },
                    "required": ["query"],
                },
            ),
            Tool(
                name="fetch_abstracts",
                description="Consulta resúmenes, autores, revistas y DOI a partir de una lista de PMID.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmids": {"type": "array", "items": {"type": "string"}}},
                    "required": ["pmids"],
                },
            ),
            Tool(
                name="find_related",
                description="Consulta artículos relacionados con un PMID mediante PubMed Similar Articles.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmid": {"type": "string"}, "top_k": {"type": "integer", "default": 10}},
                    "required": ["pmid"],
                },
            ),
            Tool(
                name="get_citations",
                description="Lista artículos que citan un PMID, basada en PubMed Central.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmid": {"type": "string"}},
                    "required": ["pmid"],
                },
            ),
            Tool(
                name="search_by_author",
                description="Busca publicaciones recientes de un autor concreto.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "author_name": {"type": "string", "description": "Por ejemplo: 'Doudna JA'"},
                        "max_results": {"type": "integer", "default": 20},
                    },
                    "required": ["author_name"],
                },
            ),
        ]

    async def dispatch_tool(self, name: str, arguments: dict) -> Any:
        dispatch_map = {
            "search_pubmed": self._search,
            "fetch_abstracts": self._fetch,
            "find_related": self._find_related,
            "get_citations": self._get_citations,
            "search_by_author": self._search_by_author,
        }
        if name not in dispatch_map:
            raise ValueError(f"Unknown tool: {name}")
        instrumented = instrument("pubmed-mcp", name)(dispatch_map[name])
        return await instrumented(**arguments)

    async def _search(
        self, query: str, max_results: int = 20,
        date_range_years: int = 5, sort: str = "relevance",
    ) -> dict:
        params = {
            "db": "pubmed",
            "term": query,
            "retmax": max_results,
            "reldate": date_range_years * 365,
            "datetype": "pdat",
            "retmode": "json",
            "sort": sort,
        }
        if self.api_key:
            params["api_key"] = self.api_key
        resp = await self.http_get_with_retry(f"{EUTILS_BASE}/esearch.fcgi", params)
        pmids = resp.json().get("esearchresult", {}).get("idlist", [])
        return {"query": query, "count": len(pmids), "pmids": pmids}

    async def _fetch(self, pmids: list[str]) -> list[dict]:
        if not pmids:
            return []
        params = {
            "db": "pubmed",
            "id": ",".join(pmids),
            "rettype": "abstract",
            "retmode": "xml",
        }
        if self.api_key:
            params["api_key"] = self.api_key
        resp = await self.http_get_with_retry(f"{EUTILS_BASE}/efetch.fcgi", params, timeout=60)
        root = ET.fromstring(resp.content)
        articles = []
        for art in root.findall(".//PubmedArticle"):
            articles.append({
                "pmid": art.findtext(".//PMID"),
                "title": art.findtext(".//ArticleTitle") or "",
                "abstract": " ".join(t.text or "" for t in art.findall(".//AbstractText")),
                "authors": [
                    f"{a.findtext('LastName') or ''} {a.findtext('Initials') or ''}".strip()
                    for a in art.findall(".//Author")
                ][:10],
                "journal": art.findtext(".//Journal/Title") or "",
                "year": art.findtext(".//PubDate/Year") or "",
                "doi": next((
                    id_.text for id_ in art.findall(".//ArticleId")
                    if id_.get("IdType") == "doi"
                ), None),
                "mesh_terms": [m.findtext(".//DescriptorName") for m in art.findall(".//MeshHeading")][:10],
            })
        return articles

    async def _find_related(self, pmid: str, top_k: int = 10) -> dict:
        params = {
            "dbfrom": "pubmed", "db": "pubmed", "id": pmid,
            "linkname": "pubmed_pubmed", "retmode": "json",
        }
        if self.api_key:
            params["api_key"] = self.api_key
        resp = await self.http_get_with_retry(f"{EUTILS_BASE}/elink.fcgi", params)
        links = resp.json().get("linksets", [{}])[0].get("linksetdbs", [])
        related = []
        for link in links:
            if link.get("linkname") == "pubmed_pubmed":
                related = [l["id"] for l in link.get("links", [])][:top_k]
                break
        return {"source_pmid": pmid, "related_pmids": related}

    async def _get_citations(self, pmid: str) -> dict:
        params = {
            "dbfrom": "pubmed", "db": "pmc", "id": pmid,
            "linkname": "pubmed_pmc_refs", "retmode": "json",
        }
        if self.api_key:
            params["api_key"] = self.api_key
        resp = await self.http_get_with_retry(f"{EUTILS_BASE}/elink.fcgi", params)
        return {"source_pmid": pmid, "cited_by": resp.json()}

    async def _search_by_author(self, author_name: str, max_results: int = 20) -> dict:
        query = f"{author_name}[Author]"
        return await self._search(query, max_results=max_results, sort="pub_date")


if __name__ == "__main__":
    server = PubMedMCPServer(api_key=os.environ.get("NCBI_API_KEY"))
    asyncio.run(server.run_stdio())

"""uniprot_mcp.py — servidor MCP para la API REST de UniProt."""

class UniProtMCPServer(BioMCPServerBase):
    """tools: search_protein · get_protein_details · get_sequence · get_features · get_orthologs · get_domains.

    UniProt REST endpoint: https://rest.uniprot.org/uniprotkb/*.
    Límite de solicitudes: ninguno; se recomiendan como máximo 20 por segundo.
    """
    # Amplía el ejemplo de UniProt de la parte 14 con features, ortólogos y dominios
    # Hereda de la clase base y define el despacho
    pass


"""pdb_mcp.py — servidor MCP para la API de RCSB PDB."""

class PDBMCPServer(BioMCPServerBase):
    """tools: search_structure · get_structure_details · fetch_pdb_file · get_ligands ·
              search_by_sequence · get_experimental_method · get_related_structures.

    RCSB PDB REST API: https://data.rcsb.org/.
    Límite de solicitudes: ninguno; se recomiendan como máximo 10 por segundo.
    """
    # Patrón similar al de la parte 14: descarga de archivos CIF y PDB, etc.
    pass


"""chembl_mcp.py — servidor MCP para la API REST de ChEMBL."""

class ChEMBLMCPServer(BioMCPServerBase):
    """tools: search_compound · get_compound_details · get_target_activity ·
              smiles_to_chembl_id · get_bioactivities · get_similar_compounds.

    ChEMBL REST: https://www.ebi.ac.uk/chembl/api/data/.
    ChEMBL es una fuente abierta estándar de datos sobre fármacos, actividad y dianas.
    """
    pass


"""blast_mcp.py — servidor MCP para NCBI BLAST QBlast."""

class BLASTMCPServer(BioMCPServerBase):
    """Herramientas: submit_blast, poll_blast, get_hits y quick_blast (integra envío, sondeo y obtención).

    QBlast requiere sondeo asíncrono. Como consume muchos recursos, el límite es estricto: una solicitud por segundo.
    """
    pass

# Dockerfile.pubmed_mcp
FROM python:3.11-slim

WORKDIR /app

# Dependencias del sistema
RUN apt-get update && apt-get install -y --no-install-recommends \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Python dependency
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Código fuente
COPY bio_mcp_base.py pubmed_mcp.py .

# Expone el puerto de métricas de Prometheus
EXPOSE 9090

# Comunicación MCP mediante stdin/stdout
CMD ["python", "pubmed_mcp.py"]

# Comprobación de estado (opcional)
HEALTHCHECK --interval=30s --timeout=10s \
  CMD curl -f http://localhost:9090/metrics || exit 1

# docker-compose.yml
version: "3.9"

services:
  pubmed-mcp:
    build:
      context: .
      dockerfile: Dockerfile.pubmed_mcp
    container_name: pubmed-mcp
    environment:
      - NCBI_API_KEY=${NCBI_API_KEY}
      - LOG_LEVEL=INFO
    ports:
      - "9091:9090"  # Métricas de Prometheus
    stdin_open: true
    tty: true
    restart: unless-stopped
    logging:
      driver: json-file
      options:
        max-size: "10m"
        max-file: "3"

  uniprot-mcp:
    build: {context: ., dockerfile: Dockerfile.uniprot_mcp}
    container_name: uniprot-mcp
    ports: ["9092:9090"]
    stdin_open: true
    tty: true
    restart: unless-stopped

  pdb-mcp:
    build: {context: ., dockerfile: Dockerfile.pdb_mcp}
    container_name: pdb-mcp
    ports: ["9093:9090"]
    stdin_open: true
    tty: true
    restart: unless-stopped

  chembl-mcp:
    build: {context: ., dockerfile: Dockerfile.chembl_mcp}
    container_name: chembl-mcp
    ports: ["9094:9090"]
    stdin_open: true
    tty: true
    restart: unless-stopped

  blast-mcp:
    build: {context: ., dockerfile: Dockerfile.blast_mcp}
    container_name: blast-mcp
    ports: ["9095:9090"]
    stdin_open: true
    tty: true
    restart: unless-stopped

  # Stack de observabilidad
  prometheus:
    image: prom/prometheus:latest
    container_name: prometheus
    volumes:
      - ./prometheus.yml:/etc/prometheus/prometheus.yml
    ports: ["9099:9090"]
    restart: unless-stopped

  grafana:
    image: grafana/grafana:latest
    container_name: grafana
    ports: ["3001:3000"]
    environment:
      - GF_SECURITY_ADMIN_PASSWORD=admin
    volumes:
      - grafana-storage:/var/lib/grafana
    restart: unless-stopped

volumes:
  grafana-storage:

"""sse_transport_wrapper.py — expone remotamente un servidor MCP mediante FastAPI y SSE."""
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from mcp.server.sse import SseServerTransport


async def sse_endpoint_factory(mcp_server_instance):
    """Expone el servidor MCP en un endpoint SSE."""
    app = FastAPI()
    transport = SseServerTransport("/messages")

    @app.get("/sse")
    async def sse_endpoint(request: Request):
        async with transport.connect_sse(request.scope, request.receive, request._send) as (read, write):
            await mcp_server_instance.server.run(
                read, write, mcp_server_instance.server.create_initialization_options(),
            )
        return StreamingResponse(iter([]), media_type="text/event-stream")

    @app.post("/messages")
    async def messages_endpoint(request: Request):
        return await transport.handle_post_message(request.scope, request.receive, request._send)

    return app


# Ejemplo de ejecución (despliegue remoto)
# from pubmed_mcp import PubMedMCPServer
# import uvicorn
# server = PubMedMCPServer(api_key=os.environ.get("NCBI_API_KEY"))
# app = await sse_endpoint_factory(server)
# uvicorn.run(app, host="0.0.0.0", port=8000)

# skills/bio-research/skill.yaml
name: bio-research
version: 1.0.0
description: |
  Agente autónomo de investigación bioinformática.
  Integra cinco servicios: PubMed, UniProt, PDB, ChEMBL y BLAST.
  Claude orquesta las herramientas de forma autónoma.
authors: ["Your Lab"]
license: MIT

mcp_servers:
  - name: pubmed
    command: docker
    args: ["run", "-i", "--rm", "-e", "NCBI_API_KEY", "bio-mcp/pubmed:latest"]
    env: ["NCBI_API_KEY"]
  - name: uniprot
    command: docker
    args: ["run", "-i", "--rm", "bio-mcp/uniprot:latest"]
  - name: pdb
    command: docker
    args: ["run", "-i", "--rm", "bio-mcp/pdb:latest"]
  - name: chembl
    command: docker
    args: ["run", "-i", "--rm", "bio-mcp/chembl:latest"]
  - name: blast
    command: docker
    args: ["run", "-i", "--rm", "bio-mcp/blast:latest"]

system_prompt: |
  Eres un asistente de investigación bioinformática.
  Llama de forma autónoma a las herramientas MCP necesarias y responde basándote en evidencia.

  Principios:
  1. No inventes hechos. Toda afirmación debe estar respaldada por resultados de herramientas.
  2. Incluye las fuentes: PMID, accession de UniProt, ID de PDB o ID de ChEMBL.
  3. Si falla una herramienta, intenta una ruta alternativa e informa del fallo.
  4. Ahorra contexto y minimiza las llamadas innecesarias.
  5. Ante un límite de solicitudes, aplica backoff automáticamente e informa al usuario de la espera.

capabilities:
  - biological literature search
  - protein sequence and structure retrieval
  - chemical compound and drug activity lookup
  - sequence similarity search (BLAST)
  - cross-reference between databases

examples:
  - "Resume las publicaciones de los últimos cinco años sobre BRCA1 humano, incluidas estructuras 3D y ligandos relacionados"
  - "Enumera proteínas diana candidatas a las que podría unirse un SMILES concreto"
  - "Busca vías, enfermedades y ensayos clínicos recientes relacionados con este gen"

"""tests/test_pubmed_mcp.py — pruebas de integración del MCP de PubMed."""
import pytest
import time
from pubmed_mcp import PubMedMCPServer


@pytest.fixture
def server():
    return PubMedMCPServer()


@pytest.mark.asyncio
async def test_search_returns_pmids(server):
    result = await server._search("BRCA1", max_results=5)
    assert result["count"] > 0
    assert all(p.isdigit() for p in result["pmids"])


@pytest.mark.asyncio
async def test_fetch_abstracts_valid_pmid(server):
    results = await server._fetch(["33746851"])  # PMID válido conocido
    assert len(results) == 1
    assert results[0]["title"]
    assert results[0]["abstract"]


@pytest.mark.asyncio
async def test_fetch_abstracts_empty_list(server):
    results = await server._fetch([])
    assert results == []


@pytest.mark.asyncio
async def test_find_related(server):
    result = await server._find_related("33746851", top_k=5)
    assert "related_pmids" in result
    assert isinstance(result["related_pmids"], list)


@pytest.mark.asyncio
async def test_search_by_author(server):
    result = await server._search_by_author("Doudna JA", max_results=5)
    assert result["count"] > 0


@pytest.mark.asyncio
async def test_rate_limit_respected(server):
    """El límite de solicitudes fuerza una espera ante llamadas rápidas consecutivas."""
    start = time.time()
    for _ in range(5):
        await server._search("test", max_results=1)
    elapsed = time.time() - start
    # Política de NCBI: 3 solicitudes/s; cinco requieren al menos ~1,3 s sin clave de API
    assert elapsed >= 1.0, f"Límite de solicitudes incumplido: {elapsed}"


@pytest.mark.asyncio
async def test_dispatch_tool_unknown(server):
    with pytest.raises(ValueError, match="Unknown tool"):
        await server.dispatch_tool("nonexistent_tool", {})


# Integración continua
# pytest tests/ -v --asyncio-mode=auto

# prometheus.yml
global:
  scrape_interval: 15s

scrape_configs:
  - job_name: 'mcp-servers'
    static_configs:
      - targets:
          - 'pubmed-mcp:9090'
          - 'uniprot-mcp:9090'
          - 'pdb-mcp:9090'
          - 'chembl-mcp:9090'
          - 'blast-mcp:9090'