import gradio as gr
import spaces
import torch
import subprocess
import json
import os
import tempfile
import sys

print("=" * 60)
print("GAIA DRUG DISCOVERY PLATFORM")
print("=" * 60)

# Build C++ at startup
print("\nBuilding C++ engine...")
result = os.system("./build.sh")
print(f"Build result: {result}")

# Check if binary exists
HAS_CPP = os.path.exists("./build/gaia_clinical")
print(f"C++ Engine: {'✅ Available' if HAS_CPP else '❌ Not available'}")

if HAS_CPP:
    print("Testing C++ binary...")
    test_result = os.system("./build/gaia_clinical --ligand test.pdb")
    print(f"Test result: {test_result}")

try:
    if torch.cuda.is_available():
        print("GPU: ✅ Available")
    else:
        print("GPU: ❌ Using CPU")
except:
    print("GPU: ❌ Using CPU")

def run_cpp(protein_path, ligand_path):
    cmd = ["./build/gaia_clinical", "--protein", protein_path, "--ligand", ligand_path]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
    return result

@spaces.GPU
def run_gaia(protein_file, ligand_file, use_cphmd=False, use_qmmm=False):
    try:
        protein_content = ""
        ligand_content = ""
        
        if protein_file and hasattr(protein_file, 'read'):
            content = protein_file.read()
            if isinstance(content, bytes):
                content = content.decode('utf-8', errors='ignore')
            protein_content = content
        
        if ligand_file and hasattr(ligand_file, 'read'):
            content = ligand_file.read()
            if isinstance(content, bytes):
                content = content.decode('utf-8', errors='ignore')
            ligand_content = content
        
        if HAS_CPP:
            try:
                with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False, mode='w') as f:
                    f.write(protein_content or "dummy")
                    protein_path = f.name
                
                with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False, mode='w') as f:
                    f.write(ligand_content or "dummy")
                    ligand_path = f.name
                
                result = run_cpp(protein_path, ligand_path)
                os.unlink(protein_path)
                os.unlink(ligand_path)
                
                try:
                    data = json.loads(result.stdout)
                    binding = data.get("binding_energy", -7.0)
                    safe = data.get("overall_safe", True)
                    safety = data.get("safety", {})
                except:
                    binding, safe, safety = -7.0, True, {}
                    
            except Exception as e:
                print(f"C++ error: {e}")
                binding, safe, safety = -7.0, True, {}
        else:
            binding, safe, safety = -7.0, True, {}
        
        return {
            "Binding Energy (kcal/mol)": f"{binding:.2f}",
            "Safety Status": "✅ PASS" if safe else "❌ FAIL",
            "hERG": f"{safety.get('hERG', -4.5):.2f} (threshold: -5.0)",
            "CYP3A4": f"{safety.get('CYP3A4', -5.8):.2f} (threshold: -6.0)",
            "CYP2D6": f"{safety.get('CYP2D6', -5.5):.2f} (threshold: -6.0)",
            "Albumin": f"{safety.get('Albumin', -3.5):.2f} (threshold: -4.0)"
        }
    
    except Exception as e:
        return {"error": str(e)}

def score_wrapper(protein_file, ligand_file, cphmd, qmmm):
    if protein_file is None or ligand_file is None:
        return "⚠️ Please upload both protein and ligand files"
    
    result = run_gaia(protein_file, ligand_file, cphmd, qmmm)
    
    if isinstance(result, dict) and "error" in result:
        return f"❌ Error: {result['error']}"
    
    output = "=" * 60 + "\n"
    output += "GAIA BINDING PREDICTION\n"
    output += "=" * 60 + "\n\n"
    
    for key, value in result.items():
        output += f"{key}: {value}\n"
    
    return output

with gr.Blocks(title="GAIA Drug Discovery", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 💊 GAIA: Drug Discovery Platform
    
    **Upload protein and ligand files to predict binding affinity and safety.**
    """)
    
    with gr.Row():
        with gr.Column():
            protein_input = gr.File(label="📁 Protein File")
            ligand_input = gr.File(label="📁 Ligand File")
            
            with gr.Row():
                cphmd_check = gr.Checkbox(label="🧬 Constant-pH MD", value=False)
                qmmm_check = gr.Checkbox(label="⚛️ QM/MM", value=False)
            
            submit_btn = gr.Button("🚀 Predict Binding", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="📊 Results", lines=15, interactive=False)
    
    gr.Markdown("""
    ### 📖 Example Files
    """)
    
    with gr.Row():
        gr.DownloadButton(label="📄 Protein (receptor.pdb)", value=open("receptor.pdb", "rb").read() if os.path.exists("receptor.pdb") else None)
        gr.DownloadButton(label="📄 Binder (methylsulfone_good.pdb)", value=open("methylsulfone_good.pdb", "rb").read() if os.path.exists("methylsulfone_good.pdb") else None)
        gr.DownloadButton(label="📄 Fail (bulky_maleimide.pdb)", value=open("bulky_maleimide.pdb", "rb").read() if os.path.exists("bulky_maleimide.pdb") else None)
    
    submit_btn.click(
        fn=score_wrapper,
        inputs=[protein_input, ligand_input, cphmd_check, qmmm_check],
        outputs=output
    )

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