"""
Download model checkpoints and install dependencies at startup.
"""
import os
import subprocess
import sys
from pathlib import Path

BASE_DIR = Path(__file__).parent


def install_pyg_deps():
    """Install PyG dependencies with correct torch version."""
    try:
        import torch
        torch_version = torch.__version__.split("+")[0]
        cuda_version = "cu" + torch.version.cuda.replace(".", "") if torch.cuda.is_available() else "cpu"
        pyg_url = f"https://data.pyg.org/whl/torch-{torch_version}+{cuda_version}.html"
        print(f"Installing PyG deps from: {pyg_url}")
        subprocess.run([
            sys.executable, "-m", "pip", "install",
            "torch-scatter", "torch-sparse", "torch-cluster", "pyg-lib",
            "-f", pyg_url, "--quiet", "--no-deps"
        ], check=False)
        print("PyG deps installed")
    except Exception as e:
        print(f"PyG deps install warning: {e}")


def download_evodiffmol_dataset():
    """Download EvoDiffMol moses dataset."""
    dataset_path = BASE_DIR / "datasets/moses/without_h/raw/valid.npz"
    if dataset_path.exists():
        print("EvoDiffMol dataset already exists")
        return
    
    print("Downloading EvoDiffMol dataset...")
    dataset_path.parent.mkdir(parents=True, exist_ok=True)
    
    try:
        from huggingface_hub import hf_hub_download
        hf_hub_download(
            repo_id="scofieldlinlin/EvoDiffMol",
            filename="datasets/moses/without_h/raw/valid.npz",
            local_dir=str(BASE_DIR),
        )
        print("EvoDiffMol dataset downloaded")
    except Exception as e:
        print(f"Dataset download warning: {e}")
        # Try alternative: download from EvoDiffMol repo scripts
        try:
            script = BASE_DIR / "EvoDiffMol/datasets/process_datasets.py"
            if script.exists():
                subprocess.run([sys.executable, str(script), "--dataset", "moses"], 
                             cwd=str(BASE_DIR / "EvoDiffMol"), check=False)
        except Exception as e2:
            print(f"Dataset process warning: {e2}")


def setup():
    # Clone RLMolLM if not exists
    if not (BASE_DIR / "RLMolLM").exists():
        print("Cloning RLMolLM...")
        subprocess.run(["git", "clone", "https://github.com/XiaoboLinlin/RLMolLM.git"],
                      cwd=BASE_DIR, check=True)

    # Clone EvoDiffMol if not exists
    if not (BASE_DIR / "EvoDiffMol").exists():
        print("Cloning EvoDiffMol...")
        subprocess.run(["git", "clone", "https://github.com/XiaoboLinlin/EvoDiffMol.git"],
                      cwd=BASE_DIR, check=True)

    # Install RLMolLM without deps (torch already installed)
    if not (BASE_DIR / "RLMolLM" / "rlmollm.egg-info").exists():
        print("Installing RLMolLM...")
        subprocess.run([sys.executable, "-m", "pip", "install", "-e",
                       str(BASE_DIR / "RLMolLM"), "--no-deps", "--quiet"], check=False)

    # Install EvoDiffMol without deps (torch already installed)
    if not (BASE_DIR / "EvoDiffMol" / "evodiffmol.egg-info").exists():
        print("Installing EvoDiffMol...")
        subprocess.run([sys.executable, "-m", "pip", "install", "-e",
                       str(BASE_DIR / "EvoDiffMol"), "--no-deps", "--quiet"], check=False)

    # Install PyG deps
    install_pyg_deps()

    # Download RLMolLM checkpoint
    rlmollm_checkpoint = BASE_DIR / "RLMolLM/assets/models/moses_pretrained.pt"
    if not rlmollm_checkpoint.exists():
        print("Downloading RLMolLM checkpoint...")
        os.chdir(BASE_DIR / "RLMolLM")
        subprocess.run([sys.executable, "download_assets.py", "--dataset", "moses"], check=False)
        os.chdir(BASE_DIR)

    # Download EvoDiffMol checkpoint
    evodiff_checkpoint = BASE_DIR / "EvoDiffMol/assets/checkpoints/moses_without_h_80.pt"
    if not evodiff_checkpoint.exists():
        print("Downloading EvoDiffMol checkpoint...")
        os.chdir(BASE_DIR / "EvoDiffMol")
        subprocess.run([sys.executable, "assets/download_checkpoint.py"], check=False)
        os.chdir(BASE_DIR)

    # Download EvoDiffMol dataset
    download_evodiffmol_dataset()

    print("Setup complete!")


if __name__ == "__main__":
    setup()