import sys

import polars as pl
import pytest
import torch

from sglang.srt.debug_utils.dump_loader import (
    LOAD_FAILED,
    ValueWithMeta,
    _add_duplicate_index,
    _cast_to_polars_dtype,
    find_row,
    parse_meta_from_filename,
    read_meta,
)
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=30, suite="default", nightly=True)


class TestReadMeta:
    def test_basic(self, tmp_path):
        for fn in [
            "step=1___rank=0___dump_index=1___name=a.pt",
            "step=2___rank=0___dump_index=2___name=b.pt",
        ]:
            torch.save(torch.randn(5), tmp_path / fn)

        df = read_meta(str(tmp_path))
        assert len(df) == 2
        assert all(c in df.columns for c in ["step", "rank", "name"])


class TestFindRow:
    def test_single_match(self):
        df = pl.DataFrame({"id": [1, 2], "name": ["a", "b"], "file": ["f1", "f2"]})
        assert find_row(df, {"id": 2})["file"] == "f2"

    def test_no_match(self):
        df = pl.DataFrame({"id": [1, 2], "name": ["a", "b"], "file": ["f1", "f2"]})
        assert find_row(df, {"id": 999}) is None

    def test_ambiguous(self):
        df = pl.DataFrame({"id": [1, 1], "file": ["f1", "f2"]})
        assert find_row(df, {"id": 1}) is None


class TestCastToPolars:
    def test_int(self):
        assert _cast_to_polars_dtype("42", pl.Int64) == 42

    def test_float(self):
        assert _cast_to_polars_dtype("3.14", pl.Float64) == pytest.approx(3.14)


class TestAddDuplicateIndex:
    def test_basic(self):
        df = pl.DataFrame(
            {
                "name": ["a", "a", "b"],
                "dump_index": [1, 2, 3],
                "filename": ["f1", "f2", "f3"],
            }
        )
        result = _add_duplicate_index(df)
        assert result.filter(pl.col("name") == "a").sort("dump_index")[
            "duplicate_index"
        ].to_list() == [0, 1]


class TestValueWithMeta:
    def test_load_dict_format(self, tmp_path) -> None:
        path = tmp_path / "step=0___rank=0___dump_index=1___name=hidden.pt"
        tensor = torch.randn(4, 8)
        torch.save({"value": tensor, "meta": {"custom": "field"}}, path)

        loaded = ValueWithMeta.load(path)
        assert torch.allclose(loaded.value, tensor)
        assert loaded.meta["custom"] == "field"
        assert loaded.meta["name"] == "hidden"
        assert loaded.meta["rank"] == 0

    def test_load_bare_tensor(self, tmp_path) -> None:
        path = tmp_path / "step=0___rank=0___dump_index=1___name=bare.pt"
        tensor = torch.randn(3, 3)
        torch.save(tensor, path)

        loaded = ValueWithMeta.load(path)
        assert torch.allclose(loaded.value, tensor)
        assert loaded.meta["name"] == "bare"

    def test_load_corrupted_file(self, tmp_path) -> None:
        path = tmp_path / "step=0___rank=0___dump_index=1___name=bad.pt"
        path.write_text("not a valid pt file")

        loaded = ValueWithMeta.load(path)
        assert loaded.value is LOAD_FAILED
        assert loaded.meta["name"] == "bad"


class TestRecomputeStatusParsing:
    def test_parse_recompute_status_from_filename(self) -> None:
        from pathlib import Path

        meta_disabled = parse_meta_from_filename(
            Path(
                "step=0___rank=0___dump_index=1___name=x___recompute_status=disabled.pt"
            )
        )
        assert meta_disabled["recompute_status"] == "disabled"

        meta_recompute = parse_meta_from_filename(
            Path(
                "step=0___rank=0___dump_index=1___name=x___recompute_status=recompute.pt"
            )
        )
        assert meta_recompute["recompute_status"] == "recompute"

        meta_original = parse_meta_from_filename(
            Path(
                "step=0___rank=0___dump_index=1___name=x___recompute_status=original.pt"
            )
        )
        assert meta_original["recompute_status"] == "original"


if __name__ == "__main__":
    sys.exit(pytest.main([__file__]))