"""LLM prompt 模块自检脚本

需先下载 MolT5 权重，默认从 models/molt5-base 读取，
可用环境变量 MOLT5_PATH 覆盖。若权重缺失则跳过（不计失败）。
"""

from __future__ import annotations

import os
import sys
import traceback
from typing import Callable, List, Tuple

import torch

from lnp_ml.modeling.layers.llm_prompt import LLMPromptEncoder

MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")
SMILES = ["CCO", "c1ccccc1"]


def _build(d_model: int = 256, **kw) -> LLMPromptEncoder:
    return LLMPromptEncoder(d_model=d_model, model_name_or_path=MOLT5_PATH, **kw)


def test_output_shape() -> None:
    enc = _build().eval()
    out = enc(SMILES)
    assert out.shape == (2, 256), f"out shape {out.shape} != (2, 256)"


def test_single_example() -> None:
    enc = _build(d_model=128).eval()
    out = enc(["CCO"])
    assert out.shape == (1, 128), f"out shape {out.shape} != (1, 128)"


def test_encoder_frozen() -> None:
    enc = _build()
    n_trainable = sum(int(p.requires_grad) for p in enc.encoder.parameters())
    assert n_trainable == 0, "默认应冻结 MolT5 encoder"
    assert all(p.requires_grad for p in enc.proj_down.parameters()), "proj_down 应可训练"


def test_cache_used() -> None:
    enc = _build().eval()
    enc(SMILES)
    assert all(s in enc._cache for s in SMILES), "冻结模式应缓存句向量"
    # 第二次前向应复用缓存（不报错且形状一致）
    assert enc(SMILES).shape == (2, 256)


def test_grad_flows_to_proj() -> None:
    """冻结 encoder 时，梯度应能流到可训练的 proj_down。"""
    enc = _build()
    enc(SMILES).sum().backward()
    g = enc.proj_down[0].weight.grad
    assert g is not None and torch.isfinite(g).all(), "proj_down 未收到有效梯度"


# ──────────────────────────────────────────────────────────────────────
# Runner
# ──────────────────────────────────────────────────────────────────────

def _collect_tests() -> List[Tuple[str, Callable[[], None]]]:
    return [
        (name, fn) for name, fn in globals().items()
        if name.startswith("test_") and callable(fn)
    ]


def main() -> int:
    try:
        _build()
    except Exception as e:
        print(f"[SKIP] 无法加载 MolT5（path={MOLT5_PATH}）：{e}")
        print("       请先下载权重，或设置 MOLT5_PATH 后重试。")
        return 0

    tests = _collect_tests()
    pass_count = 0
    fail: List[Tuple[str, str]] = []
    for name, fn in tests:
        try:
            fn()
        except Exception:
            fail.append((name, traceback.format_exc()))
            print(f"[FAIL] {name}")
        else:
            pass_count += 1
            print(f"[ OK ] {name}")

    print(f"\n{'='*60}")
    print(f"Passed: {pass_count}/{len(tests)}")
    if fail:
        print(f"\n{'='*60}\nFailures:")
        for name, tb in fail:
            print(f"\n--- {name} ---\n{tb}")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())