import json
import sys
from types import SimpleNamespace

import numpy as np
import pytest

from topovec.viewer.document import (
    load_npz_document,
    load_sad_document,
    load_viewer_document,
)


def test_load_npz_document_enumerates_states(tmp_path):
    path = tmp_path / "field.npz"
    states = np.zeros((2, 3, 4, 5, 1, 3), dtype=np.float32)
    states[1, ..., 2] = 1.0
    np.savez_compressed(
        path,
        PATH=states,
        settings=json.dumps({"L": 10}),
    )

    document = load_npz_document(path)

    assert document.format_name == "NPZ"
    assert tuple(document.system.size) == (3, 4, 5)
    assert [entry.label for entry in document.entries] == ["PATH[0]", "PATH[1]"]
    assert document.entries[1].load_state().shape == (3, 4, 5, 1, 3)
    assert np.allclose(document.entries[1].load_state()[..., 2], 1.0)


def test_load_viewer_document_rejects_unknown_suffix(tmp_path):
    path = tmp_path / "field.txt"
    path.write_text("x")

    with pytest.raises(ValueError, match="Unsupported viewer document type"):
        load_viewer_document(path)


def test_load_sad_document_uses_lazy_reader(tmp_path, monkeypatch):
    path = tmp_path / "field.sad"
    path.write_bytes(b"fake")

    class WrongReader:
        def __init__(self, filename):
            raise AssertionError("ContainerReader should not be used for viewer SAD loading")

    class FakeRandomReader:
        def __init__(self, filename):
            self.filename = filename
            self.nkeys = 2
            self._index = 0
            self.closed = False

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            self.close()

        def seek(self, index):
            self._index = int(index)

        def __getitem__(self, key):
            if key == "x":
                return np.full((2, 3, 4), self._index + 1.0)
            if key == "y":
                return np.full((2, 3, 4), self._index + 2.0)
            if key == "z":
                return np.full((2, 3, 4), self._index + 3.0)
            if key == "s":
                return {"frame": self._index}
            raise KeyError(key)

        def close(self):
            self.closed = True

    monkeypatch.setitem(
        sys.modules,
        "sadcompressor",
        SimpleNamespace(
            SADRandomReader=FakeRandomReader,
            ContainerReader=WrongReader,
        ),
    )

    document = load_sad_document(path)

    assert document.format_name == "SAD"
    assert len(document.entries) == 2
    state0 = document.entries[0].load_state()
    state1 = document.entries[1].load_state()
    assert state0.shape == (4, 3, 2, 1, 3)
    assert not np.array_equal(state0, state1)
    assert document.metadata["settings"] == {"frame": 0}
