"""bio_mcp_base.py — すべての生物MCPサーバーの共通基本クラス。"""
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:
    """レート制限ポリシー。"""
    max_requests_per_second: float = 3.0
    max_requests_per_hour: int | None = None
    burst: int = 5


class RateLimiter:
    """トークンバケットレートリミッター（非同期）。"""

    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ツール呼び出し回数",
    ["server", "tool", "status"],
)
TOOL_LATENCY = Histogram(
    "mcp_tool_latency_seconds",
    "MCPツールレイテンシー（秒）",
    ["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):
    """ツール呼び出しレイテンシーと成功率を記録するデコレーター。"""
    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サーバーの基本クラス。"""

    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標準ハンドラーを登録します。"""
        @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]:
        """このサーバーで公開されるツールのリスト。"""
        ...

    @abstractmethod
    async def dispatch_tool(self, name: str, arguments: dict) -> Any:
        """ツール名に基づいてディスパッチします。"""
        ...

    async def http_get_with_retry(
        self,
        url: str,
        params: dict | None = None,
        timeout: int = 30,
    ) -> requests.Response:
        """レート制限と再試行を行う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:  # レート制限に達しました
                    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トランスポート。"""
        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キーを使用する場合は10リクエスト/秒、使用しない場合は3リクエスト/秒。
        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="PubMedの類似論文に基づいて、PMIDに関連する論文を取得します。",
                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):
    """ツール：search_protein、get_protein_details、get_sequence、get_features、get_orthologs、get_domains。

    UniProt RESTエンドポイント：https://rest.uniprot.org/uniprotkb/。
    レート制限：なし（推奨：≤20/秒）。
    """
    # トピック14のUniProtの例を拡張（特徴、オルソログ、ドメインを追加）
    # 基本クラスを継承し、ディスパッチを定義します。
    pass


"""pdb_mcp.py — RCSB PDB API MCPサーバー。"""

class PDBMCPServer(BioMCPServerBase):
    """ツール：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/。
    レート制限：なし（推奨：≤10/秒）。
    """
    # （トピック14、CIF、PDBファイル取得などと同様のパターン）
    pass


"""chembl_mcp.py — ChEMBL REST API MCPサーバー。"""

class ChEMBLMCPServer(BioMCPServerBase):
    """ツール：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):
    """ツール：submit_blast、poll_blast、get_hits、quick_blast（送信+ポーリング+取得を統合）。

    QBlastには非同期ポーリングが必要です。リソースを大量に消費するため、レート制限は厳密です（1/秒）。
    """
    pass

# Dockerfile.pubmed_mcp
FROM python:3.11-slim

WORKDIR /app

# システム依存関係
RUN apt-get update && apt-get install -y --no-install-recommends \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Python依存関係
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# ソース
COPY bio_mcp_base.py pubmed_mcp.py .

# Prometheusメトリクスのポートを公開
EXPOSE 9090

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

# ヘルスチェック（オプション）
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 — MCPサーバーを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):
    """MCPサーバーを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


# 実行例（リモートデプロイ）
# 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がツールを自律的に調整します。
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ツールを自律的に呼び出し、証拠に基づいて回答します。

  原則：
  1. 事実を捏造しないでください。すべての主張は、ツール呼び出しの結果に基づいて行う必要があります。
  2. ソース（PMID、UniProtアクセッション番号、PDB ID、ChEMBL ID）を回答に引用してください。
  3. ツール呼び出しが失敗した場合は、別の方法を試して、失敗を報告してください。
  4. コンテキストを保存します。不必要なツール呼び出しを最小限に抑えます。
  5. レート制限の場合は、自動的に待機し、ユーザーに待機時間を通知します。

capabilities:
  - 生物学的文献の検索
  - タンパク質配列と構造の取得
  - 化学化合物と薬物の活性の検索
  - 配列類似性検索（BLAST）
  - データベース間の相互参照

examples:
  - 「ヒトのBRCA1に関する過去5年間の論文を要約し、関連する3D構造とリガンドを含めてください。」
  - 「特定のSMILESに結合する可能性のある標的タンパク質のリストを作成します。」
  - 「この遺伝子に関連する経路、疾患、最新の臨床試験をリストします。」

"""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):
    """レート制限により、連続したリクエストで待機する必要があります。"""
    start = time.time()
    for _ in range(5):
        await server._search("test", max_results=1)
    elapsed = time.time() - start
    # NCBIポリシー：3リクエスト/秒。5つのリクエストには、少なくとも約1.3秒かかるはずです（APIキーがない場合）。
    assert elapsed >= 1.0, f"レート制限が適用されていません：{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'