"""uniprot_mcp_server.py — Exposes the UniProt REST API as an MCP tool.

An MCP server can expose multiple tools. Here are three:
- search_protein: Search for proteins using a text query
- get_protein_details: Retrieve detailed information by accession
- get_sequence: Retrieve the FASTA sequence by accession
"""
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="Search for proteins in UniProt using a text query. Returns the top K results.",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Search term (e.g., 'p53 human')"},
                    "limit": {"type": "integer", "default": 10},
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="get_protein_details",
            description="Retrieve detailed information for a protein from UniProt using its accession.",
            inputSchema={
                "type": "object",
                "properties": {"accession": {"type": "string"}},
                "required": ["accession"],
            },
        ),
        Tool(
            name="get_sequence",
            description="Retrieve the FASTA sequence for a UniProt accession.",
            inputSchema={
                "type": "object",
                "properties": {"accession": {"type": "string"}},
                "required": ["accession"],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
    """Dispatches to the appropriate function based on the tool name. In a real-world scenario, it is recommended to separate each function."""
    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 failed: {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 — MCP tool for PubMed E-utilities."""
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="Search PubMed for a list of papers (PMIDs) using a search term.",
            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="Retrieve detailed information (abstracts, authors, journals) for a list of PMIDs.",
            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 (summary, only the core logic)."""
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:
    """Returns the QBlast RID. Polling is required afterward."""
    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 parsing failed")
    return m.group(1)


def poll_blast(rid: str, poll_interval: int = 30, max_wait: int = 600) -> str:
    """Waits for the BLAST result. Returns XML when complete."""
    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 failed RID={rid}")
        time.sleep(poll_interval)
    raise TimeoutError(f"BLAST timeout RID={rid}")

# The MCP tool wrapper follows the same pattern as the UniProt and PubMed servers. The code is omitted for brevity.

"""bio_agent.py — Claude Agent + multi-MCP server orchestration."""
import asyncio
from contextlib import AsyncExitStack

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


class BioMCPAgent:
    """Claude Agent that orchestrates multiple MCP servers."""

    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:
        """Connects to an MCP server. Spawns the process and retrieves the 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()
        # Converts to Claude's tool use schema
        for tool in tools_result.tools:
            self.all_tools.append({
                "name": f"{name}__{tool.name}",  # Prevents collisions by prefixing with the server 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:
        """Runs an autonomous loop. Claude repeats until 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

            # Processes the tool_use block
            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:
                    # Calls the tool using the original name after removing the server name prefix
                    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],  # Saves context
                    })
                    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 call failed: {e}",
                        "is_error": True,
                    })
            messages.append({"role": "user", "content": tool_results})

        # Extracts the final text
        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 = """You are a bioinformatics research assistant.
You understand user requests and autonomously call the necessary MCP tools (uniprot__, pubmed__, blast__)
to answer based on evidence.

Principles:
1. Do not create facts. All claims must be supported by tool call results.
2. Always include the source (UniProt accession, PMID, etc.) in the answer.
3. If a tool call fails, try alternative methods, but also inform the user of the failure.
4. Minimize unnecessary tool calls (to save context and cost).
"""


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="Summarize the main findings of the top 3 recent papers on human TP53, and compare them with the functional annotations of this protein in UniProt to see if there are any new insights.",
            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 — Set of tools for paper reproduction (summary)."""

# Tool 1: fetch_paper_pdf — Fetch PDF by DOI (prioritize open access)
# Tool 2: extract_method_section — Extract text from the method section of the PDF
# Tool 3: find_github_repo — Search for a GitHub URL in the paper body and supplementary materials
# Tool 4: clone_repo — Git clone (in a sandbox folder)
# Tool 5: setup_env — Detect and attempt to install requirements.txt or conda environment
# Tool 6: run_reproduction — Run the standard entry point of the repository (README or setup.py)

# Each tool follows the same MCP pattern as the UniProt and PubMed servers.
# For safety, cloning and execution must be done in a separate container (Docker sandbox).

# skills/bio-research/skill.yaml
name: bio-research
version: 1.0.0
description: Bioinformatics research autonomous agent. Integrates UniProt, PubMed, BLAST, and paper reproduction.
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: |
  You are a bioinformatics research assistant. ...