import json
from mcp.client.streamable_http import streamablehttp_client
from mcp import ClientSession
class DrugSDAClient:
    def __init__(self, server_url: str):
        self.server_url = server_url
        self.session = None
    async def connect(self):
        print(f"server url: {self.server_url}")
        try:
            self.transport = streamablehttp_client(
                url=self.server_url,
                headers={"SCP-HUB-API-KEY": "<your-api-key>"}
            )
            self.read, self.write, self.get_session_id = await self.transport.__aenter__()
            self.session_ctx = ClientSession(self.read, self.write)
            self.session = await self.session_ctx.__aenter__()
            await self.session.initialize()
            session_id = self.get_session_id()
            print(f"✓ connect success")
            return True
        except Exception as e:
            print(f"✗ connect failure: {e}")
            import traceback
            traceback.print_exc()
            return False
    async def disconnect(self):
        try:
            if self.session:
                await self.session_ctx.__aexit__(None, None, None)
            if hasattr(self, 'transport'):
                await self.transport.__aexit__(None, None, None)
            print("✓ already disconnect")
        except Exception as e:
            print(f"✗ disconnect error: {e}")
    def parse_result(self, result):
        try:
            if hasattr(result, 'content') and result.content:
                content = result.content[0]
                if hasattr(content, 'text'):
                    return json.loads(content.text)
            return str(result)
        except Exception as e:
            return {"error": f"parse error: {e}", "raw": str(result)}

## Initialize clients for both Tool and Model servers
tool_client = DrugSDAClient("https://scp.intern-ai.org.cn/api/v1/mcp/2/DrugSDA-Tool")
model_client = DrugSDAClient("https://scp.intern-ai.org.cn/api/v1/mcp/3/DrugSDA-Model")
if not await tool_client.connect() or not await model_client.connect():
    print("connection failed")
    return
## Input: List of candidate SMILES strings
smiles_list = ['O=C(Nc1cccc2c1CCCC2)N1CCc2c([nH]c3ccccc23)C1c1cccc(F)c1F', ...]
## Step 1: Calculate QED scores
result = await tool_client.session.call_tool(
    "calculate_mol_drug_chemistry",
    arguments={"smiles_list": smiles_list}
)
QED_result = tool_client.parse_result(result)["metrics"]
## Step 2: Predict ADMET properties (LD50)
result = await model_client.session.call_tool(
    "pred_molecule_admet",
    arguments={"smiles_list": smiles_list}
)
LD50_result = model_client.parse_result(result)["admet_preds"]
## Step 3: Filter molecules by QED and LD50 criteria
select_smiles_list = []
for i in range(len(smiles_list)):
    QED = QED_result[i]["qed"]
    LD50 = LD50_result[i]["LD50_Zhu"]
    if QED >= 0.6 and LD50 >= 3.0:
        select_smiles_list.append(smiles_list[i])
## Step 4: Retrieve protein structure by PDB code
pdb_code = "6vkv"
result = await tool_client.session.call_tool(
    "retrieve_protein_data_by_pdbcode",
    arguments={"pdb_code": pdb_code}
)
pdb_path = tool_client.parse_result(result)["pdb_path"]
## Step 5: Extract main chain
result = await tool_client.session.call_tool(
    "save_main_chain_pdb",
    arguments={"pdb_file": pdb_path, "main_chain_id": ""}
)
pdb_path = tool_client.parse_result(result)["out_file"]
## Step 6: Fix PDB file for docking
result = await tool_client.session.call_tool(
    "fix_pdb_dock",
    arguments={"pdb_file_path": pdb_path}
)
pdb_path = tool_client.parse_result(result)["fix_pdb_file_path"]
## Step 7: Identify binding pocket
result = await model_client.session.call_tool(
    "run_fpocket",
    arguments={"pdb_path": pdb_path}
)
best_pocket = tool_client.parse_result(result)["pockets"][0]
## Step 8: Convert SMILES to PDBQT format
result = await tool_client.session.call_tool(
    "convert_smiles_to_other_format",
    arguments={"inputs": select_smiles_list, "target_format": "pdbqt"}
)
ligand_paths = [x["output_file"] for x in tool_client.parse_result(result)["convert_results"]]
## Step 9: Convert protein PDB to PDBQT
result = await tool_client.session.call_tool(
    "convert_pdb_to_pdbqt_dock",
    arguments={"input_pdb_path": pdb_path}
)
receptor_path = tool_client.parse_result(result)["output_file"]
## Step 10: Perform molecular docking
result = await model_client.session.call_tool(
    "quick_molecule_docking",
    arguments={
        "receptor_path": receptor_path,
        "ligand_paths": ligand_paths,
        "center_x": best_pocket["center_x"],
        "center_y": best_pocket["center_y"],
        "center_z": best_pocket["center_z"],
        "size_x": best_pocket["size_x"],
        "size_y": best_pocket["size_y"],
        "size_z": best_pocket["size_z"]
    }
)
docking_results = model_client.parse_result(result)["docking_results"]
## Step 11: Filter by binding affinity
final_smiles_list = []
for item in docking_results:
    if item['affinity'] <= -7.0:
        final_smiles_list.append(select_smiles_list[item['index']])
print(f"Final candidates: {final_smiles_list}")
await tool_client.disconnect()
await model_client.disconnect()