"""
Example: using LiteFold Lab from the SDK.

Demonstrates:
  1. Structure prediction  – submit, poll, get results
  2. Pocket detection      – one-shot call
  3. Docking               – submit, poll, get summary
  4. SBDD                  – submit, poll, list molecules
  5. Protein design        – validate-then-submit, poll, get content

Run:
    pip install -e ../
    LITEFOLD_API_KEY=lf_... python lab_usage.py
"""

import os
import time

from litefold import Client, LiteFoldError

API_KEY = os.environ["LITEFOLD_API_KEY"]
BASE_URL = os.environ.get("LITEFOLD_BASE_URL", "https://agentsapi.litefold.ai")
PROJECT = "Default"


def poll_structure(client: Client, job_name: str, interval: int = 10):
    while True:
        status = client.lab.structure.get_status(PROJECT, job_name)
        print(f"  structure: {status.job_status}  {status.completed}/{status.total_files}")
        if status.job_status in ("completed", "failed"):
            return status
        time.sleep(interval)


def poll_docking(client: Client, job_name: str, protein: str, interval: int = 10):
    while True:
        status = client.lab.docking.get_status(PROJECT, job_name, protein)
        print(f"  docking: {status.job_status}  {status.completed}/{status.total_ligands}")
        if status.job_status in ("completed", "failed"):
            return status
        time.sleep(interval)


def poll_sbdd(client: Client, job_name: str, protein: str, interval: int = 10):
    while True:
        status = client.lab.sbdd.get_status(PROJECT, job_name, protein)
        print(f"  sbdd: {status.job_status}  generated={status.total_generated}/{status.total_requested}")
        if status.job_status in ("completed", "failed"):
            return status
        time.sleep(interval)


def poll_design(client: Client, job_name: str, model: str, config: str, interval: int = 10):
    while True:
        status = client.lab.design.get_status(PROJECT, job_name, model, config)
        print(f"  design: {status.status}")
        if status.status in ("completed", "failed"):
            return status
        time.sleep(interval)


def main():
    client = Client(api_key=API_KEY, base_url=BASE_URL)

    # ── 1. Structure Prediction ──────────────────────────────────
    print("=== Structure Prediction ===")
    sub = client.lab.structure.submit(PROJECT, "struct-run-1", ["protein.fasta"])
    print(f"Submitted: job_id={sub.job_id}  success={sub.success}")

    jobs = client.lab.structure.list_jobs(PROJECT)
    print(f"Total structure jobs: {len(jobs)}")

    status = poll_structure(client, "struct-run-1")
    if status.job_status == "completed":
        for f in status.files:
            if f.status == "completed":
                result = client.lab.structure.get_result(PROJECT, "struct-run-1", f.file_name)
                print(f"  {f.file_name}: pLDDT={result.metrics.mean_plddt if result.metrics else 'N/A'}")

    # ── 2. Pocket Detection (one-shot) ───────────────────────────
    print("\n=== Pocket Detection ===")
    pockets = client.lab.docking.find_pockets(PROJECT, "protein.pdb")
    print(f"Found {pockets.num_pockets} pocket(s)")
    for p in pockets.pockets[:3]:
        print(f"  #{p.pocket_id}  center={p.center}  score={p.score:.3f}")

    # ── 3. Docking ───────────────────────────────────────────────
    if pockets.pockets:
        print("\n=== Docking ===")
        best_pocket = pockets.pockets[0]
        dock_sub = client.lab.docking.submit(
            PROJECT, "dock-run-1", "protein.pdb",
            ligand_file_names=["lig1.sdf", "lig2.sdf"],
            pocket_coords=best_pocket.center,
            box_sizes=[20.0, 20.0, 20.0],
        )
        print(f"Docking submitted: job_id={dock_sub.job_id}")

        dock_status = poll_docking(client, "dock-run-1", "protein.pdb")
        if dock_status.job_status == "completed":
            summary = client.lab.docking.get_summary(PROJECT, "dock-run-1", "protein.pdb")
            if summary.summary:
                print(f"  Best score: {summary.summary.best_score}")
            for lig in (dock_status.ligands or []):
                if lig.status == "completed":
                    res = client.lab.docking.get_result(PROJECT, "dock-run-1", "protein.pdb", lig.ligand_name)
                    print(f"  {lig.ligand_name}: best={res.best_score}")

    # ── 4. SBDD ──────────────────────────────────────────────────
    if pockets.pockets:
        print("\n=== SBDD ===")
        sbdd_sub = client.lab.sbdd.submit(
            PROJECT, "sbdd-run-1", "protein.pdb",
            pocket_coords=pockets.pockets[0].center,
            num_mols=50,
        )
        print(f"SBDD submitted: job_id={sbdd_sub.job_id}")

        sbdd_status = poll_sbdd(client, "sbdd-run-1", "protein.pdb")
        if sbdd_status.job_status == "completed":
            results = client.lab.sbdd.get_results(PROJECT, "sbdd-run-1", "protein.pdb", limit=5)
            print(f"  Top {len(results.molecules)} molecules:")
            for mol in results.molecules:
                print(f"    {mol.id}: vina={mol.vina_score}  qed={mol.qed}  smiles={mol.smiles}")

    # ── 5. Protein Design ────────────────────────────────────────
    print("\n=== Protein Design ===")
    config_yaml = "residues_to_design: [10, 20, 30]"
    result = client.lab.design.validate_and_submit(
        PROJECT, "design-run-1", "boltzgen", "default",
        protein_name="protein.pdb",
        config_yaml_content=config_yaml,
        num_designs=5,
    )
    print(f"Design submit: success={result.success}  job_id={result.job_id}")

    if result.success:
        design_status = poll_design(client, "design-run-1", "boltzgen", "default")
        if design_status.status == "completed":
            content = client.lab.design.get_content(PROJECT, "design-run-1", "boltzgen", "default")
            print(f"  Total designs: {content.total_designs}")

    client.close()


if __name__ == "__main__":
    try:
        main()
    except LiteFoldError as e:
        print(f"LiteFold SDK error: {e}")
