import os
import json
import requests
import gradio as gr

# ========== 配置 ==========
API_URL = "https://connect.patsnap.com/096456/Logic-mcp"
API_KEY = os.getenv("PATSNAP_API_KEY", "")

HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json",
}


def call_mcp_tool(tool_name: str, arguments: dict) -> dict:
    """通用的 MCP 工具调用函数"""
    payload = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": arguments,
        },
    }
    response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
    response.raise_for_status()
    return response.json()


def search_drugs(drug_name, target, disease, phase, limit):
    """调用 ls_drug_search 并格式化结果"""
    if not API_KEY:
        return "❌ Error: PATSNAP_API_KEY not set. Please add it as a Space Secret."

    args = {}
    if drug_name.strip():
        args["drug"] = [drug_name.strip()]
    if target.strip():
        args["target"] = [target.strip()]
    if disease.strip():
        args["disease"] = [disease.strip()]
    if phase:
        args["highest_phase"] = [phase]
    args["limit"] = int(limit)

    try:
        result = call_mcp_tool("ls_drug_search", args)
        content = result.get("result", {}).get("content", [{}])
        text = content[0].get("text", json.dumps(result)) if content else json.dumps(result)
        data = json.loads(text) if isinstance(text, str) else text

        total_hits = data.get("total_hits", 0)
        results_list = data.get("results", [])

        if not results_list:
            return f"📭 No drugs found (total: {total_hits})"

        output = f"### Search Results (showing {len(results_list)} of {total_hits} total hits)\n\n"
        output += "| Drug Name | Highest Phase | Target | Disease | Organization |\n"
        output += "|-----------|---------------|--------|---------|--------------|\n"
        for row in results_list:
            name = row.get("drug_name", row.get("name", "N/A"))
            phase_val = row.get("highest_phase", "N/A")
            target_val = row.get("target", "N/A")
            disease_val = row.get("disease", "N/A")
            org = row.get("organization", "N/A")
            output += f"| {name} | {phase_val} | {target_val} | {disease_val} | {org} |\n"
        return output
    except Exception as e:
        return f"❌ Search failed: {str(e)}"


# ========== Gradio 界面 ==========
with gr.Blocks(title="PatSnap Pharma Intelligence — Drug Search") as demo:
    gr.Markdown("""
    # 💊 PatSnap Pharma Intelligence — Drug Search
    
    Search PatSnap's global pharmaceutical database for drugs matching your criteria.
    
    **How to use**:
    1. Fill in one or more search fields (drug name, target, disease, phase).
    2. Click **Search** to retrieve matching drug candidates.
    
    > 🔑 **API Key required**: Set `PATSNAP_API_KEY` as a Space Secret.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            drug_name = gr.Textbox(label="Drug Name", placeholder="e.g. pembrolizumab")
            target = gr.Textbox(label="Target", placeholder="e.g. PD-1")
            disease = gr.Textbox(label="Disease", placeholder="e.g. non-small cell lung cancer")
            phase = gr.Dropdown(
                choices=[
                    "discovery", "preclinical", "phase_1", "phase_2", "phase_3",
                    "approved", "phase_4"
                ],
                label="Highest Phase",
                value=None,
            )
            limit = gr.Slider(1, 50, value=10, step=1, label="Results Limit")
            search_btn = gr.Button("🔍 Search", variant="primary")

        with gr.Column(scale=2):
            output_md = gr.Markdown("### Results will appear here")

    search_btn.click(
        fn=search_drugs,
        inputs=[drug_name, target, disease, phase, limit],
        outputs=[output_md],
    )


if __name__ == "__main__":
    demo.launch()