"""Regression tests for the 2026-10-14 audit findings U03-3, U03-4, U05-1, U11-8, U11-9, U11-18
(LoRA export fidelity, adapter-to-base binding on resume, export-check gating and GPU pinning).

Run: cd /apps/development/rlhf/align && CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest \
        tests/test_audit_g2_lora_export_fixes.py -q -p no:cacheprovider
"""

from __future__ import annotations

import importlib.util
import logging
import os

import pytest
import torch
from test_qwen3_parity import make_tiny, rescale_

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _load(name):
    spec = importlib.util.spec_from_file_location(name, os.path.join(REPO, "scripts", f"{name}.py"))
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


exp = _load("qwen3_export_check")


# ---------------------------------------------------------------------------
# U05-1 / U03-3: the bf16 merge cast is lossy; the export must MEASURE and (optionally) refuse
# ---------------------------------------------------------------------------


def _bf16_lora(b_scale: float):
    from align.ppo.lora import LoRAConfig, inject_lora, lora_modules

    m = make_tiny().to(torch.bfloat16)
    inject_lora(m, LoRAConfig(r=4, alpha=8.0), adapter="sft", generator=torch.Generator().manual_seed(1))
    g = torch.Generator().manual_seed(2)
    with torch.no_grad():
        for _, mod in lora_modules(m):
            mod.adapters["sft"].B.copy_(b_scale * torch.randn(mod.adapters["sft"].B.shape, generator=g))
    return m


def test_export_reports_retained_delta_and_warns_when_bf16_drops_most_of_it(caplog):
    from align.ppo.lora import export_merged_state_dict

    m = _bf16_lora(1e-5)  # the shipped PPO lr=1e-6 regime: delta << half an ulp of W
    rep: dict = {}
    with caplog.at_level(logging.WARNING):
        export_merged_state_dict(m, {k: v.detach().clone() for k, v in m.state_dict().items()}, "sft", report=rep)
    assert rep["cosine"] < 0.5 and rep["dropped_fraction"] > 0.5, rep
    assert rep["layers"] > 0 and rep["min_layer_cosine"] <= rep["cosine"] + 1e-6
    assert any("bf16" in r.message and "cosine" in r.message for r in caplog.records)


def test_export_refuses_below_min_cosine_and_passes_a_well_resolved_delta(caplog):
    from align.ppo.lora import export_merged_state_dict

    full = lambda m: {k: v.detach().clone() for k, v in m.state_dict().items()}  # noqa: E731
    small = _bf16_lora(1e-5)
    with pytest.raises(ValueError, match="cosine"):
        export_merged_state_dict(small, full(small), "sft", min_cosine=0.9)
    big = _bf16_lora(5.0)  # delta >> ulp: the cast preserves it
    rep: dict = {}
    with caplog.at_level(logging.WARNING):
        export_merged_state_dict(big, full(big), "sft", report=rep, min_cosine=0.9)
    assert rep["cosine"] > 0.95
    assert not any("cosine" in r.message for r in caplog.records)


def test_export_of_an_exactly_zero_adapter_is_reported_not_nan():
    from align.ppo.lora import export_merged_state_dict

    m = _bf16_lora(0.0)
    rep: dict = {}
    export_merged_state_dict(m, {k: v.detach().clone() for k, v in m.state_dict().items()}, "sft", report=rep)
    assert rep["delta_norm"] == 0.0 and rep["cosine"] == 1.0


# ---------------------------------------------------------------------------
# U03-4: a LoRA resume binds the adapter to the base it was trained on
# ---------------------------------------------------------------------------


def test_lora_payload_carries_a_base_fingerprint_and_a_foreign_base_is_refused(tmp_workdir, monkeypatch):
    import test_qwen3_stage as T

    tr, _, _ = T._build(tmp_workdir, "a")
    payload = tr.train_module.optim_state_dict()
    assert isinstance(payload.get("base_fingerprint"), str) and len(payload["base_fingerprint"]) == 64
    # the same base resumes fine
    tr_same, _, _ = T._build(tmp_workdir, "same")
    tr_same.train_module.load_optim_state_dict(payload)
    # a different base with an identical key set / shapes (repointed snapshot) must not
    monkeypatch.setattr(T, "make_tiny", lambda: rescale_(make_tiny(), seed=7))
    tr_other, _, _ = T._build(tmp_workdir, "other")
    with pytest.raises(RuntimeError, match="base"):
        tr_other.train_module.load_optim_state_dict(payload)


def test_legacy_payload_without_a_fingerprint_still_resumes_with_a_warning(tmp_workdir, caplog):
    import test_qwen3_stage as T

    tr, _, _ = T._build(tmp_workdir, "a")
    payload = tr.train_module.optim_state_dict()
    payload.pop("base_fingerprint")
    with caplog.at_level(logging.WARNING):
        tr.train_module.load_optim_state_dict(payload)
    assert any("fingerprint" in r.message for r in caplog.records)


def test_fingerprint_tracks_content_not_just_keys():
    from align.train.qwen3_stage import base_fingerprint

    a, b = make_tiny(), make_tiny()
    assert base_fingerprint(a.state_dict()) == base_fingerprint(b.state_dict())
    with torch.no_grad():
        b.model.layers[0].attn.q_norm.weight.mul_(1.001)
    assert base_fingerprint(a.state_dict()) != base_fingerprint(b.state_dict())
    other = rescale_(make_tiny(), seed=7)
    assert base_fingerprint(a.state_dict()) != base_fingerprint(other.state_dict())


# ---------------------------------------------------------------------------
# U11-8: the exit status gates every check the script advertises
# ---------------------------------------------------------------------------


def _good_report():
    return {
        "verdict": {"passed": True},
        "merge_vs_trainer_file": {"bitwise_mismatches": 0},
        "snapshot_mapping_roundtrip": {"bitwise_equal": True},
        "merge_changed_tensors_vs_snapshot": 12,
        "merge_relative_weight_change": {"n": 12, "min": 1e-4, "mean": 4e-4, "max": 1e-3},
        "adapter_vs_merged": {"verdict": {"passed": True}},
    }


def test_export_ok_passes_a_complete_report():
    assert exp.export_ok(_good_report()) is True


@pytest.mark.parametrize("mutate", [
    lambda r: r.update(snapshot_mapping_roundtrip={"bitwise_equal": False}),
    lambda r: r.update(merge_changed_tensors_vs_snapshot=0),                       # zero adapter
    lambda r: r.update(merge_relative_weight_change={"n": 12, "min": 0.0, "mean": 1e-4, "max": 1e-3}),
    lambda r: r.pop("merge_changed_tensors_vs_snapshot"),                          # check did not run
    lambda r: r.pop("snapshot_mapping_roundtrip"),
    lambda r: r.update(adapter_vs_merged={"verdict": {"passed": False}}),
    lambda r: r.pop("adapter_vs_merged"),
    lambda r: r.update(merge_vs_trainer_file={"bitwise_mismatches": 2}),
    lambda r: r.update(verdict={"passed": False}),
    lambda r: r.pop("verdict"),
])
def test_export_ok_fails_each_unmet_check(mutate):
    r = _good_report()
    mutate(r)
    assert exp.export_ok(r) is False


# ---------------------------------------------------------------------------
# U11-9: merged-bf16 vs adapter-applied forward
# ---------------------------------------------------------------------------


def _lp(t=64, v=50, seed=0):
    g = torch.Generator().manual_seed(seed)
    return torch.log_softmax(3 * torch.randn(t, v, generator=g), -1)


def test_adapter_vs_merged_passes_identical_and_fails_a_dropped_delta():
    a = _lp()
    same = exp.evaluate_adapter_vs_merged(exp.compare_adapter_vs_merged([a], [a.clone()]))
    assert same["passed"] is True
    other = _lp(seed=5)  # an unrelated distribution: what a merge that lost the delta looks like
    bad = exp.evaluate_adapter_vs_merged(exp.compare_adapter_vs_merged([a], [other]))
    assert bad["passed"] is False
    near = a + 0.02 * torch.randn(a.shape, generator=torch.Generator().manual_seed(3))
    ok = exp.evaluate_adapter_vs_merged(exp.compare_adapter_vs_merged([a], [torch.log_softmax(near, -1)]))
    assert ok["passed"] is True


def test_adapter_vs_merged_shape_mismatch_raises():
    with pytest.raises(ValueError):
        exp.compare_adapter_vs_merged([_lp(8)], [_lp(9)])


def test_merge_fidelity_stats_expose_the_dead_zone():
    from align.ppo.lora import lora_state_dict

    m = _bf16_lora(1e-5)
    base = {k: v.detach().clone() for k, v in m.state_dict().items() if ".adapters." not in k}
    lora = lora_state_dict(m, "sft")
    merged = exp.merge_lora_state(base, lora, 2.0)
    st = exp.merge_fidelity_stats(base, merged, lora, 2.0)
    assert st["cosine"] < 0.5 and st["dropped_fraction"] > 0.5 and st["layers"] > 0


# ---------------------------------------------------------------------------
# U11-18: the vLLM child is pinned to the PARENT's first visible device
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("cvd,expect", [("2,3", "2"), ("3", "3"), (" 1 , 2", "1"), ("GPU-abc,GPU-def", "GPU-abc"),
                                        (None, "0"), ("", "0")])
def test_first_visible_gpu_follows_the_parent_pinning(cvd, expect):
    env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd}
    assert exp._first_visible_gpu(env) == expect


def test_rollout_env_uses_the_pinned_device_not_physical_zero(monkeypatch):
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3")
    assert exp._rollout_env(exp._first_visible_gpu())["CUDA_VISIBLE_DEVICES"] == "2"