"""推理 profiler 脚本：CPU 上也要能跑出耗时结果（有 GPU 时额外给出显存峰值）。"""

import argparse
import importlib.util
from pathlib import Path

import pytest

torch = pytest.importorskip("torch")

from models import build_model  # noqa: E402

SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "profile_inference.py"


@pytest.fixture(scope="module")
def script():
    spec = importlib.util.spec_from_file_location("profile_inference", SCRIPT)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


@pytest.fixture
def model():
    return build_model(
        {
            "name": "defect_net",
            "num_classes": 2,
            "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
            "neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
            "head": {"hidden_channels": 8},
        }
    ).eval()


@pytest.fixture
def args():
    return argparse.Namespace(iters=2, warmup=1, amp=False, postprocess=True, rows=3, trace=None)


def test_measure_reports_latency_and_memory(script, model, args):
    images = torch.randn(2, 3, 32, 32)
    result = script.measure(model, images, args)

    assert result["batch_size"] == 2
    assert result["latency_ms"] > 0
    assert result["throughput_fps"] > 0
    if torch.cuda.is_available():
        assert result["peak_allocated_mb"] > 0
        assert result["peak_reserved_mb"] >= result["peak_allocated_mb"]
    else:
        assert "peak_allocated_mb" not in result


def test_profile_exports_trace(script, model, args, tmp_path, capsys):
    args.trace = str(tmp_path / "trace.json")
    script.profile_once(model, torch.randn(1, 3, 32, 32), args)

    assert Path(args.trace).exists()
    out = capsys.readouterr().out
    assert "显存占用 Top" in out and "耗时 Top" in out


def test_postprocess_toggle_runs_both_paths(script, model, args):
    images = torch.randn(1, 3, 32, 32)
    args.postprocess = False
    script.run_once(model, images, args)
    args.postprocess = True
    script.run_once(model, images, args)