"""Standalone unit tests for the TRELLIS-HiCache patch logic.

No ComfyUI, no GPU, no TRELLIS: a dummy DiT and a fake SparseTensor stand in for
the real models, so the scheduling / CFG-routing / run-boundary / copy-on-patch
logic is checked deterministically on CPU.
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import pytest
from trellis_hicache_patch import (
    HiCacheModelPatch, apply_hicache, remove_hicache, validate_config,
)


class DummyDiT(torch.nn.Module):
    """Counts forward calls; returns a deterministic dense velocity."""
    def __init__(self, dim=8):
        super().__init__()
        self.dim = dim
        self.calls = 0
        self.lin = torch.nn.Linear(dim, dim)

    def forward(self, x, t, cond=None, **kw):
        self.calls += 1
        # smooth-in-t signal so the Hermite forecast is accurate
        tv = float(t.reshape(-1)[0]) if torch.is_tensor(t) else float(t)
        return torch.ones(1, self.dim) * tv


class FakeSparse:
    """Minimal stand-in for a TRELLIS SparseTensor: .feats + .replace()."""
    def __init__(self, feats):
        self.feats = feats
    def replace(self, feats):
        return FakeSparse(feats)


def _trellis_t_seq(steps=25):
    # TRELLIS scales t to [0, 1000], decreasing.
    import numpy as np
    return [float(x) for x in np.linspace(1, 0, steps + 1)[:-1] * 1000.0]


def test_validate_config_rejects_bad_params():
    with pytest.raises(ValueError):
        validate_config("nope", 3, 2, 1, 0.5, 5)
    with pytest.raises(ValueError):
        validate_config("hermite", 0, 2, 1, 0.5, 5)
    with pytest.raises(ValueError):
        validate_config("hermite", 3, 2, 1, 1.5, 5)


def test_skips_steps_on_decreasing_trellis_schedule():
    dit = DummyDiT()
    patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=2)
    for t in _trellis_t_seq(25):
        patch(torch.zeros(1, 8), t)
    # 25 steps, interval 3: clearly fewer real forwards than steps, some skipped.
    assert patch.skipped_steps > 0
    assert dit.calls == patch.computed_steps
    assert patch.computed_steps + patch.skipped_steps == 25


def test_split_cfg_routes_two_states():
    """Each timestep called twice (cond, uncond) must use two forecast states and
    still skip; the DiT must be called far fewer than 2x steps."""
    dit = DummyDiT()
    patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=2)
    for t in _trellis_t_seq(25):
        patch(torch.zeros(1, 8), t)   # cond
        patch(torch.zeros(1, 8), t)   # uncond (repeated t)
    assert patch.skipped_steps > 0
    assert patch.computed_steps + patch.skipped_steps == 50
    assert dit.calls < 50            # genuinely skipped DiT forwards


def test_new_run_resets_on_direction_reversal():
    dit = DummyDiT()
    patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=2)
    for t in _trellis_t_seq(10):
        patch(torch.zeros(1, 8), t)
    first = patch.computed_steps + patch.skipped_steps
    # a second run restarts t near 1000 (direction flips up) -> reset
    for t in _trellis_t_seq(10):
        patch(torch.zeros(1, 8), t)
    assert patch.computed_steps + patch.skipped_steps == 10  # counters reset, not 20
    assert first == 10


def test_sparse_output_forecast_rebuilds_sparse():
    """SLaT-style sparse output: forecast must return a FakeSparse, not a tensor."""
    class SparseDiT(torch.nn.Module):
        def __init__(self): super().__init__(); self.calls = 0
        def forward(self, x, t, cond=None, **kw):
            self.calls += 1
            tv = float(t.reshape(-1)[0]) if torch.is_tensor(t) else float(t)
            return FakeSparse(torch.ones(4, 8) * tv)
    patch = HiCacheModelPatch(SparseDiT(), method="hermite", interval=2, warmup_steps=1)
    outs = [patch(None, t) for t in _trellis_t_seq(12)]
    assert all(isinstance(o, FakeSparse) for o in outs)
    assert patch.skipped_steps > 0


def test_apply_remove_is_copy_on_patch():
    class Pipe:
        def __init__(self):
            self.models = {"sparse_structure_flow_model": DummyDiT(),
                           "shape_slat_flow_model_512": DummyDiT(), "shape_slat_flow_model_1024": DummyDiT()}
    p = Pipe()
    patched = apply_hicache(p, method="hermite", interval=2, stages="both")
    # original untouched
    assert not getattr(p.models["shape_slat_flow_model_512"], "_hicache_is_patch", False)
    assert getattr(patched.models["shape_slat_flow_model_512"], "_hicache_is_patch", False)
    assert patched.models is not p.models
    # remove restores
    clean = remove_hicache(patched)
    assert not getattr(clean.models["shape_slat_flow_model_512"], "_hicache_is_patch", False)


def test_apply_never_nests():
    class Pipe:
        def __init__(self):
            self.models = {"sparse_structure_flow_model": DummyDiT(),
                           "shape_slat_flow_model_512": DummyDiT(), "shape_slat_flow_model_1024": DummyDiT()}
    p = Pipe()
    a = apply_hicache(p, interval=2, stages="both")
    b = apply_hicache(a, interval=3, stages="both")  # re-patch
    inner = b.models["shape_slat_flow_model_512"].inner
    assert not getattr(inner, "_hicache_is_patch", False)  # unwrapped, not nested
    assert b.models["shape_slat_flow_model_512"].interval == 3


def test_wrap_none_does_not_crash_on_attr_access():
    """Regression for issue #1: GGUF/lazy pipelines leave pipeline.models[key]=None
    at patch time. Wrapping None must not raise, and an unknown-attribute lookup
    must give a clear 'not loaded yet' error, NOT the misleading
    'HiCacheModelPatch object has no attribute inner'."""
    patch = HiCacheModelPatch(None, method="hermite", interval=3, warmup_steps=2)
    assert patch.inner is None                     # readable, no crash
    with pytest.raises(AttributeError) as ei:
        _ = patch.some_config_attr                 # forwarded lookup on a None inner
    msg = str(ei.value)
    assert "not loaded yet" in msg and "some_config_attr" in msg
    # a compute step on an unbound patch fails loudly, not silently
    with pytest.raises(RuntimeError):
        patch(torch.zeros(1, 8), _trellis_t_seq(4)[0])


def test_bind_inner_materializes_lazy_model():
    """After the real DiT loads, bind_inner attaches it and forwards run/skip."""
    patch = HiCacheModelPatch(None, method="hermite", interval=3, warmup_steps=2)
    dit = DummyDiT()
    assert patch.bind_inner(dit) is patch
    assert patch.inner is dit
    assert "inner" in patch._modules          # registered as a submodule
    for t in _trellis_t_seq(25):
        patch(torch.zeros(1, 8), t)
    assert patch.skipped_steps > 0
    assert dit.calls == patch.computed_steps
    assert patch.computed_steps + patch.skipped_steps == 25


def test_lazy_pipeline_guard_then_load_still_triggers():
    """A lazy loader's ``is None`` guard must still run, including after unload."""
    class LazyPipe:
        def __init__(self):
            self.models = {"sparse_structure_flow_model": None,
                           "shape_slat_flow_model_512": None,
                           "shape_slat_flow_model_1024": None}
            self.loads = []

        def load_sparse_structure_model(self):
            if self.models["sparse_structure_flow_model"] is None:
                model = DummyDiT()
                self.loads.append(model)
                self.models["sparse_structure_flow_model"] = model
            return self.models["sparse_structure_flow_model"]

    p = LazyPipe()
    patched = apply_hicache(p, method="hermite", interval=3, warmup_steps=2, stages="sparse_structure")
    assert patched.models["sparse_structure_flow_model"] is None

    slot = patched.load_sparse_structure_model()
    assert getattr(slot, "_hicache_is_patch", False)
    assert slot.inner is patched.loads[0]
    for t in _trellis_t_seq(25):
        slot(torch.zeros(1, 8), t)
    assert slot.skipped_steps > 0
    assert patched.loads[0].calls == slot.computed_steps

    patched.models["sparse_structure_flow_model"] = None
    assert patched.models["sparse_structure_flow_model"] is None
    replacement = patched.load_sparse_structure_model()
    assert len(patched.loads) == 2
    assert replacement is not slot
    assert getattr(replacement, "_hicache_is_patch", False)
    assert replacement.inner is patched.loads[1]


def test_remove_hicache_clears_pending_before_any_load():
    """Disabling a pending patch must not re-enable it on a later assignment."""
    class LazyPipe:
        def __init__(self):
            self.models = {"sparse_structure_flow_model": None,
                           "shape_slat_flow_model_512": None,
                           "shape_slat_flow_model_1024": None}

    patched = apply_hicache(
        LazyPipe(), method="hermite", interval=3, warmup_steps=2,
        stages="sparse_structure",
    )
    clean = remove_hicache(patched)
    assert type(clean.models) is dict
    assert clean.models["sparse_structure_flow_model"] is None

    dit = DummyDiT()
    clean.models["sparse_structure_flow_model"] = dit
    assert clean.models["sparse_structure_flow_model"] is dit


def test_eager_inner_still_registered_as_submodule():
    """Eager path unbroken: a real nn.Module inner stays a registered submodule so
    state_dict / parameters / device moves recurse into it."""
    dit = DummyDiT()
    patch = HiCacheModelPatch(dit, method="hermite", interval=3)
    assert patch.inner is dit
    assert "inner" in patch._modules
    # inner's parameters are reachable through the patch (state_dict recurses)
    assert any(k.startswith("inner.") for k in patch.state_dict().keys())
    assert any(p is q for q in dit.parameters() for p in patch.parameters())


def test_stages_selector():
    class Pipe:
        def __init__(self):
            self.models = {"sparse_structure_flow_model": DummyDiT(),
                           "shape_slat_flow_model_512": DummyDiT(), "shape_slat_flow_model_1024": DummyDiT()}
    p = Pipe()
    ss = apply_hicache(p, stages="sparse_structure")
    assert getattr(ss.models["sparse_structure_flow_model"], "_hicache_is_patch", False)
    assert not getattr(ss.models["shape_slat_flow_model_512"], "_hicache_is_patch", False)