"""uniprot_mcp_server.py — UniProt REST APIをMCPツールとして公開します。

MCPサーバーは、複数のツールを公開できます。以下に3つの例を示します。
- search_protein：テキストクエリを使用してタンパク質を検索します。
- get_protein_details：アクセッション番号に基づいて詳細情報を取得します。
- get_sequence：アクセッション番号に基づいてFASTA形式の配列を取得します。
"""
import asyncio
import json
from typing import Any

import requests
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent


UNIPROT_BASE = "https://rest.uniprot.org"
server = Server("uniprot-mcp")


@server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="search_protein",
            description="UniProtでテキストクエリを使用してタンパク質を検索します。上位K件の結果を返します。",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "検索用語（例：'p53 human'）"},
                    "limit": {"type": "integer", "default": 10},
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="get_protein_details",
            description="UniProtのアクセッション番号を使用して、タンパク質の詳細情報を取得します。",
            inputSchema={
                "type": "object",
                "properties": {"accession": {"type": "string"}},
                "required": ["accession"],
            },
        ),
        Tool(
            name="get_sequence",
            description="UniProtアクセッション番号のFASTA配列を取得します。",
            inputSchema={
                "type": "object",
                "properties": {"accession": {"type": "string"}},
                "required": ["accession"],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
    """ツール名に基づいて適切な関数にディスパッチします。実用的なシナリオでは、各関数を分離することをお勧めします。"""
    try:
        if name == "search_protein":
            result = _search_protein(arguments["query"], arguments.get("limit", 10))
        elif name == "get_protein_details":
            result = _get_details(arguments["accession"])
        elif name == "get_sequence":
            result = _get_sequence(arguments["accession"])
        else:
            raise ValueError(f"不明なツール：{name}")
        return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False, indent=2))]
    except requests.RequestException as e:
        return [TextContent(type="text", text=json.dumps({"error": f"UniProt APIが失敗しました：{e}"}))]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({"error": str(e)}))]


def _search_protein(query: str, limit: int) -> dict:
    resp = requests.get(
        f"{UNIPROT_BASE}/uniprotkb/search",
        params={
            "query": query,
            "format": "json",
            "size": limit,
            "fields": "accession,id,protein_name,organism_name,length",
        },
        timeout=30,
    )
    resp.raise_for_status()
    results = resp.json().get("results", [])
    return {"count": len(results), "results": results}


def _get_details(accession: str) -> dict:
    resp = requests.get(
        f"{UNIPROT_BASE}/uniprotkb/{accession}.json",
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    return {
        "accession": data.get("primaryAccession"),
        "protein_name": data.get("proteinDescription", {}),
        "organism": data.get("organism", {}),
        "length": data.get("sequence", {}).get("length"),
        "function": [
            c["texts"][0]["value"]
            for c in data.get("comments", [])
            if c.get("commentType") == "FUNCTION"
        ],
    }


def _get_sequence(accession: str) -> dict:
    resp = requests.get(f"{UNIPROT_BASE}/uniprotkb/{accession}.fasta", timeout=30)
    resp.raise_for_status()
    lines = resp.text.strip().split("\n")
    return {
        "accession": accession,
        "header": lines[0],
        "sequence": "".join(lines[1:]),
    }


async def main():
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


if __name__ == "__main__":
    asyncio.run(main())

"""pubmed_mcp_server.py — PubMed E-utilities用のMCPツール。"""
import asyncio
import json
from xml.etree import ElementTree as ET

import requests
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent


EUTILS_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
server = Server("pubmed-mcp")


@server.list_tools()
async def list_tools() -> 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},
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="fetch_abstracts",
            description="PMIDのリストの詳細情報（抄録、著者、ジャーナル）を取得します。",
            inputSchema={
                "type": "object",
                "properties": {
                    "pmids": {"type": "array", "items": {"type": "string"}},
                },
                "required": ["pmids"],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    try:
        if name == "search_pubmed":
            result = _search(arguments["query"], arguments.get("max_results", 20), arguments.get("date_range_years", 5))
        elif name == "fetch_abstracts":
            result = _fetch(arguments["pmids"])
        else:
            raise ValueError(f"不明なツール：{name}")
        return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False, indent=2))]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({"error": str(e)}))]


def _search(query: str, max_results: int, date_range_years: int) -> dict:
    esearch = requests.get(
        f"{EUTILS_BASE}/esearch.fcgi",
        params={
            "db": "pubmed",
            "term": query,
            "retmax": max_results,
            "reldate": date_range_years * 365,
            "datetype": "pdat",
            "retmode": "json",
        },
        timeout=30,
    )
    esearch.raise_for_status()
    pmids = esearch.json().get("esearchresult", {}).get("idlist", [])
    return {"query": query, "count": len(pmids), "pmids": pmids}


def _fetch(pmids: list[str]) -> list[dict]:
    if not pmids:
        return []
    efetch = requests.get(
        f"{EUTILS_BASE}/efetch.fcgi",
        params={"db": "pubmed", "id": ",".join(pmids), "rettype": "abstract", "retmode": "xml"},
        timeout=60,
    )
    efetch.raise_for_status()
    root = ET.fromstring(efetch.content)
    articles = []
    for art in root.findall(".//PubmedArticle"):
        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")
        ]
        journal = art.findtext(".//Journal/Title") or ""
        year = art.findtext(".//PubDate/Year") or ""
        articles.append({
            "pmid": pmid,
            "title": title,
            "abstract": abstract,
            "authors": authors[:10],
            "journal": journal,
            "year": year,
        })
    return articles


async def main():
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


if __name__ == "__main__":
    asyncio.run(main())

"""blast_mcp_server.py — NCBI BLAST QBlast（概要、コアロジックのみ）。"""
import time
import re
import requests

BLAST_URL = "https://blast.ncbi.nlm.nih.gov/Blast.cgi"


def submit_blast(sequence: str, program: str = "blastp", database: str = "nr") -> str:
    """QBlast RIDを返します。後でポーリングする必要があります。"""
    resp = requests.post(BLAST_URL, data={
        "CMD": "Put",
        "PROGRAM": program,
        "DATABASE": database,
        "QUERY": sequence,
    }, timeout=60)
    resp.raise_for_status()
    m = re.search(r"RID = (\S+)", resp.text)
    if not m:
        raise RuntimeError("QBlast RIDの解析に失敗しました")
    return m.group(1)


def poll_blast(rid: str, poll_interval: int = 30, max_wait: int = 600) -> str:
    """BLASTの結果を待ちます。完了したらXMLを返します。"""
    start = time.time()
    while time.time() - start < max_wait:
        resp = requests.get(BLAST_URL, params={"CMD": "Get", "RID": rid, "FORMAT_OBJECT": "SearchInfo"}, timeout=30)
        if "Status=READY" in resp.text:
            xml_resp = requests.get(BLAST_URL, params={"CMD": "Get", "RID": rid, "FORMAT_TYPE": "XML"}, timeout=60)
            xml_resp.raise_for_status()
            return xml_resp.text
        elif "Status=FAILED" in resp.text:
            raise RuntimeError(f"BLASTが失敗しました RID={rid}")
        time.sleep(poll_interval)
    raise TimeoutError(f"BLASTタイムアウト RID={rid}")

# MCPツールラッパーは、UniProtとPubMedサーバーと同じパターンに従います。コードは簡潔にするために省略されています。

"""bio_agent.py — Claude Agent + 複数のMCPサーバーのオーケストレーション。"""
import asyncio
from contextlib import AsyncExitStack

import anthropic
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


class BioMCPAgent:
    """複数のMCPサーバーをオーケストレーションするClaude Agent。"""

    def __init__(self, model: str = "claude-sonnet-4-5", max_turns: int = 20):
        self.anthropic = anthropic.Anthropic()
        self.model = model
        self.max_turns = max_turns
        self.sessions: dict[str, ClientSession] = {}
        self.all_tools: list[dict] = []
        self._exit_stack: AsyncExitStack | None = None

    async def connect_server(self, name: str, command: str, args: list[str]) -> None:
        """MCPサーバーに接続します。プロセスを起動し、ツールリストを取得します。"""
        params = StdioServerParameters(command=command, args=args)
        stdio_transport = await self._exit_stack.enter_async_context(stdio_client(params))
        session = await self._exit_stack.enter_async_context(ClientSession(*stdio_transport))
        await session.initialize()
        tools_result = await session.list_tools()
        # Claudeのツール使用スキーマに変換します
        for tool in tools_result.tools:
            self.all_tools.append({
                "name": f"{name}__{tool.name}",  # サーバー名のプレフィックスを付けて競合を防ぎます
                "description": tool.description,
                "input_schema": tool.inputSchema,
            })
            self.sessions[f"{name}__{tool.name}"] = session

    async def __aenter__(self):
        self._exit_stack = AsyncExitStack()
        await self._exit_stack.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)

    async def run(self, user_query: str, system_prompt: str) -> dict:
        """自律的なループを実行します。Claudeは、stop_reason="end_turn"になるまで繰り返します。"""
        messages = [{"role": "user", "content": user_query}]
        tool_log = []

        for turn in range(self.max_turns):
            response = self.anthropic.messages.create(
                model=self.model,
                max_tokens=4096,
                system=system_prompt,
                tools=self.all_tools,
                messages=messages,
            )
            messages.append({"role": "assistant", "content": response.content})

            if response.stop_reason == "end_turn":
                break
            if response.stop_reason != "tool_use":
                continue

            # tool_useブロックを処理します
            tool_results = []
            for block in response.content:
                if block.type != "tool_use":
                    continue
                session = self.sessions.get(block.name)
                if session is None:
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": f"不明なツール：{block.name}",
                        "is_error": True,
                    })
                    continue
                try:
                    # サーバー名のプレフィックスを削除した元の名前を使用してツールを呼び出します
                    original_name = block.name.split("__", 1)[1]
                    result = await session.call_tool(original_name, block.input)
                    content_text = "\n".join(
                        c.text for c in result.content if hasattr(c, "text")
                    )
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": content_text[:10000],  # コンテキストを保存します
                    })
                    tool_log.append({
                        "tool": block.name,
                        "input": block.input,
                        "output_preview": content_text[:500],
                    })
                except Exception as e:
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": f"ツール呼び出しに失敗しました：{e}",
                        "is_error": True,
                    })
            messages.append({"role": "user", "content": tool_results})

        # 最終的なテキストを抽出します
        final_text = ""
        for block in response.content:
            if hasattr(block, "text"):
                final_text += block.text
        return {"final_answer": final_text, "tool_log": tool_log, "turns": turn + 1}


SYSTEM_PROMPT = """あなたはバイオインフォマティクスの研究アシスタントです。
ユーザーのリクエストを理解し、必要なMCPツール（uniprot__、pubmed__、blast__）を自律的に呼び出して、証拠に基づいて回答します。

原則：
1. 事実を捏造しないでください。すべての主張は、ツール呼び出しの結果によって裏付けられている必要があります。
2. 常に、回答にソース（UniProtアクセッション番号、PMIDなど）を含めてください。
3. ツールの呼び出しが失敗した場合は、代替方法を試みますが、ユーザーに失敗を通知します。
4. 不要なツール呼び出しを最小限に抑えます（コンテキストとコストを節約するため）。
"""


async def main():
    async with BioMCPAgent() as agent:
        await agent.connect_server("uniprot", "python", ["uniprot_mcp_server.py"])
        await agent.connect_server("pubmed", "python", ["pubmed_mcp_server.py"])
        # await agent.connect_server("blast", "python", ["blast_mcp_server.py"])

        result = await agent.run(
            user_query="最近の人間のTP53に関する上位3つの論文を要約し、それらをUniProtのこのタンパク質の機能注釈と比較して、新しい洞察があるかどうかを確認してください。",
            system_prompt=SYSTEM_PROMPT,
        )
        print("=== 最終的な回答 ===")
        print(result["final_answer"])
        print(f"\n=== ターン数：{result['turns']} ===")
        print(f"=== ツール呼び出し回数：{len(result['tool_log'])} ===")
        for entry in result["tool_log"]:
            print(f"  - {entry['tool']}: {entry['input']}")


if __name__ == "__main__":
    asyncio.run(main())

"""paper_repro_mcp_server.py — 論文の再現のためのツールのセット（概要）。"""

# ツール1：fetch_paper_pdf — DOIによるPDFの取得（オープンアクセスを優先します）
# ツール2：extract_method_section — PDFから方法セクションのテキストを抽出します
# ツール3：find_github_repo — 論文本文と補足資料からGitHub URLを検索します
# ツール4：clone_repo — Gitクローン（サンドボックスフォルダー内）
# ツール5：setup_env — requirements.txtまたはconda環境を検出し、インストールを試みます
# ツール6：run_reproduction — リポジトリの標準のエントリポイント（READMEまたはsetup.py）を実行します

# 各ツールは、UniProtとPubMedサーバーと同じMCPパターンに従います。
# 安全のため、クローンと実行は、別のコンテナー（Dockerサンドボックス）内で行う必要があります。

# skills/bio-research/skill.yaml
name: bio-research
version: 1.0.0
description: バイオインフォマティクス研究のための自律型エージェント。UniProt、PubMed、BLAST、および論文の再現を統合します。
mcp_servers:
  - name: uniprot
    command: python
    args: [/opt/mcp/uniprot_mcp_server.py]
  - name: pubmed
    command: python
    args: [/opt/mcp/pubmed_mcp_server.py]
  - name: blast
    command: python
    args: [/opt/mcp/blast_mcp_server.py]
  - name: paper_repro
    command: python
    args: [/opt/mcp/paper_repro_mcp_server.py]
system_prompt: |
  あなたはバイオインフォマティクス研究のアシスタントです。...