"""bio_mcp_base.py — 모든 바이오 MCP 서버의 공통 base class."""
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 정책."""
    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 메트릭 (모든 서버 공유)
TOOL_CALLS = Counter(
    "mcp_tool_calls_total",
    "MCP tool 호출 수",
    ["server", "tool", "status"],
)
TOOL_LATENCY = Histogram(
    "mcp_tool_latency_seconds",
    "MCP tool latency (초)",
    ["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):
    """tool 호출 latency · 성공률 로깅 decorator."""
    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):
    """모든 바이오 MCP 서버의 base class."""

    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  # 포트 이미 사용 중이면 스킵
        self._register_handlers()

    def _register_handlers(self):
        """MCP 표준 handler 등록."""
        @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 요청 실패: {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]:
        """이 서버가 노출하는 tool 리스트."""
        ...

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

    async def http_get_with_retry(
        self,
        url: str,
        params: dict | None = None,
        timeout: int = 30,
    ) -> requests.Response:
        """Rate limit + retry가 적용된 GET 요청."""
        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"최대 재시도 초과: {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 서버 (프로덕션급)."""
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):
        # NCBI API key가 있으면 초당 10 req 가능, 없으면 3 req.
        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="PubMed 검색어로 논문 PMID 리스트 반환. 날짜 범위·정렬 지정 가능.",
                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="PMID 리스트로 초록·저자·저널·DOI 상세 조회.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmids": {"type": "array", "items": {"type": "string"}}},
                    "required": ["pmids"],
                },
            ),
            Tool(
                name="find_related",
                description="PMID의 관련 논문 (PubMed similar articles) 조회.",
                inputSchema={
                    "type": "object",
                    "properties": {"pmid": {"type": "string"}, "top_k": {"type": "integer", "default": 10}},
                    "required": ["pmid"],
                },
            ),
            Tool(
                name="get_citations",
                description="PMID를 인용한 논문 리스트 (PubMed Central 기반).",
                inputSchema={
                    "type": "object",
                    "properties": {"pmid": {"type": "string"}},
                    "required": ["pmid"],
                },
            ),
            Tool(
                name="search_by_author",
                description="특정 저자의 최근 논문 검색.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "author_name": {"type": "string", "description": "예: '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 서버."""

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: 없음 (권장 초당 20 이하).
    """
    # 편 14 UniProt 예시 확장 (feature · ortholog · domain 추가)
    # base class 상속 · dispatch 정의
    pass


"""pdb_mcp.py — RCSB PDB API MCP 서버."""

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: 없음 (권장 초당 10 이하).
    """
    # (편 14 유사 패턴 · CIF · PDB 파일 fetch 등)
    pass


"""chembl_mcp.py — ChEMBL REST API MCP 서버."""

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은 약물 · 활성 · 표적 데이터의 표준 오픈 소스.
    """
    pass


"""blast_mcp.py — NCBI BLAST QBlast MCP 서버."""

class BLASTMCPServer(BioMCPServerBase):
    """tools: submit_blast · poll_blast · get_hits · quick_blast (submit+poll+get 통합).
    
    QBlast는 async 폴링 필요. 리소스 큰 태스크라 rate limit 엄격 (초당 1).
    """
    pass

# Dockerfile.pubmed_mcp
FROM python:3.11-slim

WORKDIR /app

# 시스템 dependency
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

# 소스
COPY bio_mcp_base.py pubmed_mcp.py .

# Prometheus 메트릭 포트 노출
EXPOSE 9090

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

# Health check (선택)
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 메트릭
    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

  # 관측성 스택
  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 — FastAPI + SSE로 MCP 서버 원격 노출."""
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from mcp.server.sse import SseServerTransport


async def sse_endpoint_factory(mcp_server_instance):
    """MCP 서버를 SSE endpoint로 노출."""
    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


# 실행 예 (원격 배포)
# 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: |
  생명정보학 리서치 자율 에이전트.
  PubMed · UniProt · PDB · ChEMBL · BLAST 5개 서비스 통합.
  Claude가 자율적으로 tool 오케스트레이션.
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: |
  당신은 생명정보학 리서치 어시스턴트입니다.
  주어진 태스크에 필요한 MCP tool을 자율적으로 호출해 증거 기반으로 답변합니다.

  원칙:
  1. 사실 창작 X. 모든 주장은 tool 호출 결과로 뒷받침.
  2. 답변에 출처 (PMID · UniProt accession · PDB ID · ChEMBL ID) 병기.
  3. tool 호출 실패 시 대안 경로 시도 + 실패 보고.
  4. 컨텍스트 절약: 불필요한 tool 호출 최소화.
  5. rate limit 시 backoff 자동, 사용자에게 대기 안내.

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

examples:
  - "human BRCA1의 최근 5년 논문 요약, 관련 3D 구조·리간드까지"
  - "특정 SMILES가 결합할 가능성 있는 표적 단백질 후보 리스트"
  - "이 유전자와 관련된 pathway · disease · 최신 임상 시험"

"""tests/test_pubmed_mcp.py — PubMed MCP 통합 테스트."""
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
    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이 대기를 강제."""
    start = time.time()
    for _ in range(5):
        await server._search("test", max_results=1)
    elapsed = time.time() - start
    # NCBI 정책 3 req/s → 5개는 최소 ~1.3초 (no api_key 가정)
    assert elapsed >= 1.0, f"Rate limit 미준수: {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 통합
# 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'