"""bio_mcp_base.py — Common base class for all bio MCP servers."""
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:
    """Rate limit policy."""
    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


# Prometheus metrics (shared across all servers)
TOOL_CALLS = Counter(
    "mcp_tool_calls_total",
    "MCP tool call count",
    ["server", "tool", "status"],
)
TOOL_LATENCY = Histogram(
    "mcp_tool_latency_seconds",
    "MCP tool latency (seconds)",
    ["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):
    """Decorator that logs tool-call latency · success rate."""
    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):
    """Base class for all bio MCP servers."""

    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  # skip if port already in use
        self._register_handlers()

    def _register_handlers(self):
        """Register MCP standard handlers."""
        @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"API request failed: {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]:
        """List of tools exposed by this server."""
        ...

    @abstractmethod
    async def dispatch_tool(self, name: str, arguments: dict) -> Any:
        """Dispatch by tool name."""
        ...

    async def http_get_with_retry(
        self,
        url: str,
        params: dict | None = None,
        timeout: int = 30,
    ) -> requests.Response:
        """GET request with rate limit + retry."""
        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"Max retries exceeded: {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 — PubMed E-utilities MCP server (production grade)."""
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):
        # 10 req/s with NCBI API key, 3 req/s without.
        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="Return PMID list for PubMed search terms. Date range · sort configurable.",
                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="Fetch abstract · authors · journal · DOI details for a PMID list.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmids": {"type": "array", "items": {"type": "string"}}},
                    "required": ["pmids"],
                },
            ),
            Tool(
                name="find_related",
                description="Fetch related articles (PubMed similar articles) for a PMID.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmid": {"type": "string"}, "top_k": {"type": "integer", "default": 10}},
                    "required": ["pmid"],
                },
            ),
            Tool(
                name="get_citations",
                description="List of articles citing this PMID (based on PubMed Central).",
                inputSchema={
                    "type": "object",
                    "properties": {"pmid": {"type": "string"}},
                    "required": ["pmid"],
                },
            ),
            Tool(
                name="search_by_author",
                description="Search recent articles by a specific author.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "author_name": {"type": "string", "description": "e.g., '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 — UniProt REST MCP server."""

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/*.
    Rate limit: none (recommend ≤ 20/s).
    """
    # Extends Topic 14 UniProt example (adds feature · ortholog · domain)
    # Inherit base class · define dispatch
    pass


"""pdb_mcp.py — RCSB PDB API MCP server."""

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/.
    Rate limit: none (recommend ≤ 10/s).
    """
    # (Similar pattern to Topic 14 · CIF · PDB file fetch etc.)
    pass


"""chembl_mcp.py — ChEMBL REST API MCP server."""

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 is the standard open source for drug · activity · target data.
    """
    pass


"""blast_mcp.py — NCBI BLAST QBlast MCP server."""

class BLASTMCPServer(BioMCPServerBase):
    """tools: submit_blast · poll_blast · get_hits · quick_blast (submit+poll+get integrated).
    
    QBlast requires async polling. It is a resource-heavy task so rate limit is strict (1/s).
    """
    pass

# Dockerfile.pubmed_mcp
FROM python:3.11-slim

WORKDIR /app

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

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

# Source
COPY bio_mcp_base.py pubmed_mcp.py .

# Expose Prometheus metrics port
EXPOSE 9090

# MCP stdin/stdout communication
CMD ["python", "pubmed_mcp.py"]

# Health check (optional)
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"  # Prometheus metrics
    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

  # Observability stack
  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 — Expose an MCP server remotely with FastAPI + 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):
    """Expose an MCP server as SSE endpoints."""
    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


# Execution example (remote deployment)
# 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: |
  Autonomous bioinformatics research agent.
  Integrates 5 services: PubMed · UniProt · PDB · ChEMBL · BLAST.
  Claude orchestrates tools autonomously.
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: |
  You are a bioinformatics research assistant.
  You autonomously call the MCP tools required for the given task and answer based on evidence.

  Principles:
  1. Do not fabricate facts. All claims must be backed by tool call results.
  2. Cite sources (PMID · UniProt accession · PDB ID · ChEMBL ID) in the answer.
  3. On tool call failure, try alternative paths + report the failure.
  4. Save context: minimize unnecessary tool calls.
  5. On rate limits, auto backoff and inform the user of the wait.

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

examples:
  - "Summarize the last 5 years of papers on human BRCA1, and include related 3D structures and ligands."
  - "List candidate target proteins that a given SMILES might bind."
  - "Pathways · diseases · latest clinical trials related to this gene."

"""tests/test_pubmed_mcp.py — PubMed MCP integration tests."""
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"])  # known valid PMID
    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):
    """Rate limit forces waiting on rapid consecutive requests."""
    start = time.time()
    for _ in range(5):
        await server._search("test", max_results=1)
    elapsed = time.time() - start
    # NCBI policy 3 req/s → 5 requests take at least ~1.3 s (assuming no api_key)
    assert elapsed >= 1.0, f"Rate limit not respected: {elapsed}"


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


# CI integration
# 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'