#!/usr/bin/env python3
"""Fetch the third-party components that are not redistributed with this release.

DegradeMaster and EGFR-PROPK are published without a license, so we cannot ship
copies of them. This script clones each upstream repository at the pinned commit
used for the paper and copies the subset the pipeline imports into the layout
`scripts/evaluate_linkers.py` and `src/classifier/protac_classifier.py` expect.

Both are optional. Generation (`scripts/generate.py`) runs without either:
  - DegradeMaster is needed only for predicted-degradation scoring.
  - EGFR-PROPK is needed only for the EGFR-only PK columns of batch_report.py.

Usage:
    python scripts/fetch_third_party.py              # fetch both
    python scripts/fetch_third_party.py degrademaster
"""
from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

COMPONENTS = {
    "degrademaster": {
        "url": "https://github.com/ABILiLab/DegradeMaster",
        "commit": "aa149beaf067a051e070b3b281f7dba43c2f3e90",
        "dest": ROOT / "src" / "degrademaster",
        "paths": [
            "model.py", "dataset.py", "nn_utils.py", "pdb_utils.py",
            "tools.py", "protacloader.py",
            "tokenizer", "utils",
            "checkpoint/1000/model_for_case_study.pth",
        ],
    },
    "propk": {
        "url": "https://github.com/Zhang-Ran-0119/EGFR-PROPK",
        "commit": "6940cb90c449e6c6d0678508fc2f228da796e028",
        "dest": ROOT / "src" / "propk",
        "paths": [
            "maplight.py", "maplight_gnn.py",
            "models/CL_1_finetune.pkl",
            "models/t0.5_10_finetune.pkl",
            "models/Vss_1_finetune.pkl",
        ],
    },
}


def fetch(name: str, spec: dict) -> int:
    print(f"\n=== {name} ===")
    print(f"  {spec['url']} @ {spec['commit'][:12]}")
    print("  NOTE: published without a license file. Review the upstream terms")
    print("        before using or redistributing it.")

    with tempfile.TemporaryDirectory() as tmp:
        try:
            subprocess.run(["git", "clone", "--quiet", spec["url"], tmp], check=True)
            subprocess.run(["git", "-C", tmp, "checkout", "--quiet", spec["commit"]],
                           check=True)
        except subprocess.CalledProcessError as exc:
            print(f"  FAILED to clone: {exc}")
            return 1

        missing = []
        for rel in spec["paths"]:
            src = Path(tmp) / rel
            dst = spec["dest"] / rel
            if not src.exists():
                missing.append(rel)
                continue
            dst.parent.mkdir(parents=True, exist_ok=True)
            if src.is_dir():
                shutil.copytree(src, dst, dirs_exist_ok=True)
            else:
                shutil.copy2(src, dst)
            print(f"  + {rel}")

        if missing:
            print(f"  {len(missing)} path(s) not found upstream at this commit:")
            for rel in missing:
                print(f"    - {rel}")
            print("  The upstream layout may have changed, or the checkpoint may be")
            print("  distributed outside git. Fetch those by hand.")
            return 1
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("components", nargs="*", default=[],
                    choices=list(COMPONENTS),
                    help="which to fetch (default: all)")
    args = ap.parse_args()

    names = args.components or list(COMPONENTS)
    failures = sum(fetch(n, COMPONENTS[n]) for n in names)

    print()
    if failures:
        print(f"{failures} component(s) incomplete.")
    else:
        print("All components fetched.")
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())