import gradio as gr

import os
import subprocess
from pathlib import Path
import shutil
import spaces

# Download model weights (skips already-downloaded models automatically)
# In total, ~6GB (3GB for RFD3, 3GB for RF3, <100MB for MPNN); may take a few minutes depending on your connection speed

cmd =  "foundry install rfd3 ligandmpnn rf3"
print("Global PATH:", os.environ.get("PATH"))
result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
if result.returncode == 0:
    print("Models installed successfully.")
else:   
    print(f"Error installing models: {result.stderr}")

#download_dir = "./models/"
#if not os.path.exists(download_dir):
#    cmd  = "foundry install rfd3 ligandmpnn rf3"
#    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
#
## Run once on startup: Install models if missing
#checkpoint_dir = Path.home() / ".foundry" / "checkpoints"
#os.environ["FOUNDRY_CHECKPOINT_DIRS"] = str(checkpoint_dir)
#
#def install_models():
#    """Download rfd3, ligandmpnn, rf3 weights once."""
#    #models = ["rfd3", "ligandmpnn", "rf3"]
#    models = ["ligandmpnn"] # let's start with only ligand mpnn for testing
#    for model in models:
#        if not (checkpoint_dir / model).exists():
#            print(f"Installing {model}...")
#            subprocess.check_call(["foundry", "install", model])
#    print("All models installed.")
#
#install_models()  # Executes on app.py load

@spaces.GPU(duration=300)  
def test_rfd3():
    """Run a quick rfd3 test design (minimal monomer, 1 step)."""
    try:
        cmd = [
            "rfd3",
            "design",
            "--seed", "42",
            "contigmap.contigs=[A10]",  # Tiny 10-res monomer
            "--num_designs", "1",
            "inference.output_prefix=test_output",
            "--inference.num_diffusion_steps", "10"  # Fast test
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode == 0:
            return "RFD3 test passed! Check test_output.pdb. Logs:\n" + result.stdout[-500:]
        else:
            return f"RFD3 test failed: {result.stderr}"
    except Exception as e:
        return f"Error: {str(e)}"
    
    
# Gradio UI
with gr.Blocks(title="RFD3 Test") as demo:
    gr.Markdown("# RFdiffusion3 (RFD3) Model Checker")
    gr.Markdown("Models auto-downloaded on launch. Click to test.")
    
    test_btn = gr.Button("Run RFD3 Test")
    output = gr.Textbox(label="Test Result")
    
    test_btn.click(test_rfd3, outputs=output)

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