"""Download and run the non-antibody catalog using pinned model, prep and Diffusers revisions."""

import argparse
import hashlib
import importlib.metadata
import json
from pathlib import Path


def pipeline_options(entry, scheduler):
    options = dict(entry["pipeline_kwargs"])
    fraction = entry.get("remaining_schedule_fraction")
    if fraction is not None:
        steps = scheduler.num_inference_steps
        scheduler.set_timesteps(steps)
        index = round((1 - fraction) * steps)
        options["partial_t"] = {
            name: float(scheduler.modality_timestep(name, index))
            for name in scheduler.modalities
        }
    return options


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo", default="nvidia/RFD4-Proteina-prep")
    parser.add_argument("--revision", default="main")
    parser.add_argument("--tasks", nargs="+")
    parser.add_argument("--seed", default=42, type=int)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--out", type=Path, default=Path("rfd4-designs"))
    parser.add_argument(
        "--prepare-only",
        action="store_true",
        help="Validate inputs on CPU without model weights.",
    )
    args = parser.parse_args()

    import torch
    from diffusers import ModularPipeline, RFD4ProteinaPipeline
    from huggingface_hub import HfApi, hf_hub_download

    revision = HfApi().model_info(args.repo, revision=args.revision).sha
    manifest_path = hf_hub_download(
        args.repo, "examples/manifest.json", revision=revision
    )
    manifest = json.loads(Path(manifest_path).read_text())
    catalog = {entry["id"]: entry for entry in manifest["examples"]}
    tasks = args.tasks if args.tasks is not None else list(catalog)
    unknown = set(tasks) - catalog.keys()
    if unknown:
        parser.error(
            f"Unknown or excluded tasks: {sorted(unknown)}. Choose from {list(catalog)}."
        )
    if len(tasks) != len(set(tasks)):
        parser.error(
            "Duplicate tasks would overwrite the same output; select each task once."
        )
    args.out.mkdir(parents=True, exist_ok=True)
    report_path = args.out / "report.json"
    destinations = [report_path] + (
        [] if args.prepare_only else [args.out / f"{name}.cif" for name in tasks]
    )
    for path in destinations:
        if path.exists() or path.is_symlink():
            raise FileExistsError(
                f"Refusing to overwrite {path}; choose a new --out directory."
            )

    helper = ModularPipeline.from_pretrained(
        args.repo, revision=revision, trust_remote_code=True
    )
    device = "cpu" if args.prepare_only else args.device
    pipe = None
    if not args.prepare_only:
        # Check the source revision before downloading weights; recent PR commits rename their keys.
        dist = importlib.metadata.distribution("diffusers")
        direct_url = json.loads(dist.read_text("direct_url.json") or "{}")
        commit = direct_url.get("vcs_info", {}).get("commit_id")
        if commit != manifest["diffusers_revision"]:
            raise RuntimeError(
                "This model snapshot requires a pinned Diffusers install. Run: pip install "
                f"git+{manifest['diffusers_repository']}@{manifest['diffusers_revision']}"
            )
        dtype = torch.bfloat16 if torch.device(device).type == "cuda" else torch.float32
        pipe = RFD4ProteinaPipeline.from_pretrained(
            manifest["model_repo"], revision=manifest["model_revision"], dtype=dtype
        ).to(device)

    report = {
        "helper_revision": revision,
        "model_revision": manifest["model_revision"],
        "seed": args.seed,
        "tasks": {},
    }
    for name in tasks:
        entry = catalog[name]
        input_path = hf_hub_download(args.repo, entry["file"], revision=revision)
        if hashlib.sha256(Path(input_path).read_bytes()).hexdigest() != entry["sha256"]:
            raise ValueError(f"Input digest mismatch: {name}")
        prepared = helper(cif_path=input_path, seed=args.seed, device=device)
        features, atoms = prepared.get("features"), prepared.get("atom_array")
        result = {
            "input_sha256": entry["sha256"],
            "tokens": int(features["padding_mask"].shape[1]),
        }
        if pipe is not None:
            options = pipeline_options(entry, pipe.scheduler)
            output = pipe(
                features,
                generator=torch.Generator(device).manual_seed(args.seed),
                **options,
            )
            path = helper(
                pipeline_output=output,
                features=features,
                atom_array=atoms,
                output_path=str(args.out / f"{name}.cif"),
                output="cif_path",
            )
            result.update(output=path, options=options)
        report["tasks"][name] = result
        print(name, result, flush=True)
    with report_path.open("x") as handle:
        json.dump(report, handle, indent=2)
        handle.write("\n")


if __name__ == "__main__":
    main()