"""uniprot_mcp_server.py — UniProt REST API를 MCP tool로 노출.

MCP 서버 하나가 tool 여러 개를 노출 가능. 여기서는 3개:
- search_protein: 텍스트 쿼리로 단백질 검색
- get_protein_details: accession으로 상세 정보
- get_sequence: accession으로 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 accession으로 단백질 상세 정보 조회.",
            inputSchema={
                "type": "object",
                "properties": {"accession": {"type": "string"}},
                "required": ["accession"],
            },
        ),
        Tool(
            name="get_sequence",
            description="UniProt accession의 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]:
    """tool 이름 dispatch. 실전에서는 각 함수 분리 권장."""
    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"Unknown tool: {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 tool."""
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"Unknown tool: {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 timeout RID={rid}")

# MCP tool wrapper는 위 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 서버 하나 연결. 프로세스 spawn + tool list 조회."""
        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 tool use 스키마로 변환
        for tool in tools_result.tools:
            self.all_tools.append({
                "name": f"{name}__{tool.name}",  # 서버 이름 prefix로 충돌 방지
                "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:
        """자율 loop. 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"Unknown tool: {block.name}",
                        "is_error": True,
                    })
                    continue
                try:
                    # MCP 서버 이름 prefix 제거 후 원본 tool 이름으로 호출
                    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"Tool 호출 실패: {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 tool(uniprot__*, pubmed__*, blast__*)을 자율적으로 호출해
증거 기반으로 답변합니다.

원칙:
1. 사실을 창작하지 않습니다. 모든 주장은 tool 호출 결과로 뒷받침되어야 합니다.
2. 답변에는 반드시 출처(UniProt accession, PMID 등)를 병기합니다.
3. tool 호출이 실패하면 대안 경로를 시도하되 사용자에게 실패 사실도 보고합니다.
4. 불필요한 tool 호출을 최소화합니다 (컨텍스트·비용 절약).
"""


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="Human TP53의 최근 5년 논문 상위 3편의 주요 발견을 요약하고, "
                       "이 단백질의 UniProt 기능 주석과 비교해 새로운 통찰이 있는지 알려주세요.",
            system_prompt=SYSTEM_PROMPT,
        )
        print("=== Final Answer ===")
        print(result["final_answer"])
        print(f"\n=== Turns: {result['turns']} ===")
        print(f"=== Tool calls: {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 — 논문 재현 tool 세트 (요약)."""

# tool 1: fetch_paper_pdf — DOI로 PDF fetch (오픈 액세스 우선)
# tool 2: extract_method_section — PDF에서 method 섹션 텍스트 추출
# tool 3: find_github_repo — 논문 본문·supplementary에서 GitHub URL 검색
# tool 4: clone_repo — git clone (샌드박스 폴더에)
# tool 5: setup_env — requirements.txt · conda env 자동 감지 · 설치 시도
# tool 6: run_reproduction — 저장소의 표준 진입점(README·setup.py) 실행

# 각 tool은 위 UniProt·PubMed 서버와 동일 MCP 패턴.
# 안전을 위해 clone/실행은 반드시 격리 컨테이너 안 (Docker sandbox).

# 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: |
  당신은 생명정보학 리서치 어시스턴트입니다. ...