"""CPU integration checks for the Hub helper; no model weights or GPU required."""

import hashlib
import importlib.util
import inspect
import json
import random
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
import torch

# `diffusers` is not a dependency of this package -- the `test` extra is pytest alone -- so this
# module is skipped rather than erroring when it is absent, matching `test_prepare.py`'s use of
# importorskip for `atomworks`. Without it, `pytest` fails at collection for anyone who installed
# this package and its documented test extra and nothing else.
pytest.importorskip("diffusers")

from diffusers import ModularPipeline, RFD4ProteinaPipeline  # noqa: E402
from rfd4_proteina_prep import prepare_cif  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
MANIFEST = json.loads((ROOT / "examples/manifest.json").read_text())


@pytest.fixture(scope="module")
def helper():
    return ModularPipeline.from_pretrained(str(ROOT), trust_remote_code=True)


@pytest.mark.parametrize("entry", MANIFEST["examples"], ids=lambda item: item["id"])
def test_preparation_matches_package_and_preserves_rng(helper, entry):
    path = ROOT / entry["file"]
    assert hashlib.sha256(path.read_bytes()).hexdigest() == entry["sha256"]
    expected, _ = prepare_cif(path, seed=42, device="cpu")
    python_state, numpy_state, torch_state = (
        random.getstate(),
        np.random.get_state(),
        torch.get_rng_state(),
    )
    prepared = helper(cif_path=str(path), seed=42, device="cpu")
    actual = prepared.get("features")
    assert random.getstate() == python_state
    now = np.random.get_state()
    assert now[0] == numpy_state[0] and now[2:] == numpy_state[2:]
    np.testing.assert_array_equal(now[1], numpy_state[1])
    assert torch.equal(torch_state, torch.get_rng_state())
    assert actual.keys() == expected.keys()
    for key in expected:
        torch.testing.assert_close(
            actual[key], expected[key], rtol=0, atol=0, equal_nan=True
        )
    inspect.signature(RFD4ProteinaPipeline.__call__).bind(None, actual)
    if entry["task"] in {"inverse_folding", "partial_diffusion"}:
        assert torch.isfinite(actual["coords_nm"][actual["atom37_mask"]]).all()


def test_catalog_excludes_antibodies_and_records_all_sources():
    ids = {entry["id"] for entry in MANIFEST["examples"]}
    assert ids == {
        "protein_binder",
        "enzyme_design",
        "inverse_folding",
        "motif_scaffolding",
        "partial_diffusion",
        "small_molecule_binder",
        "unconditional",
        "ppi_pdl1",
        "mcsa_ph2",
        "disco_iai",
    }
    for entry in MANIFEST["examples"]:
        assert entry["task"] != "antibody_design"
        assert not any(
            word in entry["file"].lower() for word in ("antibody", "nanobody", "vhh")
        )
        assert len(entry["source_revision"]) == 40
        assert (ROOT / entry["file"]).stat().st_size == entry["size_bytes"]


def test_requires_trust_for_custom_code():
    with pytest.raises(ValueError, match="trust_remote_code"):
        ModularPipeline.from_pretrained(str(ROOT), trust_remote_code=False)


@pytest.mark.parametrize(
    "kwargs,error",
    [({"seed": True}, TypeError), ({"path_length_cutoff": 0}, ValueError)],
)
def test_invalid_preparation(helper, kwargs, error):
    with pytest.raises(error):
        helper(cif_path=str(ROOT / "examples/studio/ppi_pdl1.cif"), **kwargs)


def test_missing_input(helper, tmp_path):
    with pytest.raises(FileNotFoundError):
        helper(cif_path=str(tmp_path / "missing.cif"))


def test_export_round_trip_and_no_overwrite(helper, tmp_path):
    # An existing structure is a fixture for IO, not a claim of model generation.
    prepared = helper(
        cif_path=str(ROOT / "examples/inputs/inverse_folding.cif"), seed=42
    )
    features, atoms = prepared.get("features"), prepared.get("atom_array")
    n = features["padding_mask"].shape[1]
    output = SimpleNamespace(
        coords_nm=features["coords_nm"],
        residue_type=features["residue_type"],
        atom_mask=features["atom37_mask"],
        insertion_logits=torch.zeros(1, n, n),
    )
    kwargs = {"pipeline_output": output, "features": features, "atom_array": atoms}
    path = tmp_path / "round-trip.cif"
    assert helper(**kwargs, output_path=str(path), output="cif_path") == str(path)
    assert path.is_file() and path.stat().st_size > 1000
    reread = helper(cif_path=str(path), seed=42).get("features")
    assert reread["padding_mask"].shape == features["padding_mask"].shape
    torch.testing.assert_close(reread["residue_type"], features["residue_type"])
    torch.testing.assert_close(
        reread["coords_nm"], features["coords_nm"], atol=1e-3, rtol=0, equal_nan=True
    )
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    with pytest.raises(FileExistsError):
        helper(**kwargs, output_path=str(path))
    assert hashlib.sha256(path.read_bytes()).hexdigest() == digest
    with pytest.raises(ValueError, match="row"):
        helper(**kwargs, output_path=str(tmp_path / "invalid.cif"), row=-1)
    with pytest.raises(ValueError, match="end in .cif"):
        helper(**kwargs, output_path=str(tmp_path / "invalid.pdb"))
    with pytest.raises(FileNotFoundError):
        helper(**kwargs, output_path=str(tmp_path / "absent" / "out.cif"))


def test_selected_batch_row(helper, tmp_path):
    prepared = helper(cif_path=str(ROOT / "examples/inputs/inverse_folding.cif"))
    features = {
        k: v.repeat(2, *([1] * (v.ndim - 1)))
        for k, v in prepared.get("features").items()
    }
    n = features["padding_mask"].shape[1]
    output = SimpleNamespace(
        coords_nm=features["coords_nm"],
        residue_type=features["residue_type"],
        atom_mask=features["atom37_mask"],
        insertion_logits=torch.zeros(2, n, n),
    )
    path = tmp_path / "second.cif"
    helper(
        pipeline_output=output,
        features=features,
        atom_array=prepared.get("atom_array"),
        output_path=str(path),
        row=1,
    )
    assert path.is_file()


def test_partial_diffusion_recipe():
    spec = importlib.util.spec_from_file_location(
        "runner", ROOT / "examples/run_hub_examples.py"
    )
    runner = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(runner)
    scheduler = SimpleNamespace(
        num_inference_steps=60,
        modalities=("bb_ca", "local_latents"),
        set_timesteps=lambda steps: None,
        modality_timestep=lambda name, index: index / 60,
    )
    entry = next(e for e in MANIFEST["examples"] if e["id"] == "partial_diffusion")
    assert runner.pipeline_options(entry, scheduler) == {
        "partial_t": {"bb_ca": 0.6, "local_latents": 0.6}
    }