"""
Unit tests for the segmentation v2 pipeline.

Covers:
- Postprocess: label_fixes, extract_segments, merge_begin_tags
- H5 loader: matrix_to_rotation_6d
- Output: save_as_json, save_as_elan, save_as_pt, save_predictions dispatcher
- Runner: InferenceRunner constructor, predict with synthetic input (mocked model)

Run with:
    pytest tests/test_segmentation.py -v
"""

from __future__ import annotations

import json
import xml.etree.ElementTree as ET
from pathlib import Path
from unittest.mock import MagicMock, patch

import numpy as np
import pytest

torch = pytest.importorskip("torch")


# ===========================================================================
# Postprocess tests
# ===========================================================================


class TestLabelFixes:
    """Test label_fixes two-pass BIO correction."""

    def test_empty_tensor(self):
        from sltk.segmentation.postprocess import label_fixes

        labels = torch.tensor([], dtype=torch.long)
        result = label_fixes(labels)
        assert len(result) == 0

    def test_all_out(self):
        from sltk.segmentation.postprocess import label_fixes

        labels = torch.tensor([0, 0, 0, 0, 0])
        result = label_fixes(labels)
        assert result.tolist() == [0, 0, 0, 0, 0]

    def test_inserts_begin_for_long_in_run(self):
        """An IN-run of length >= min_segment_len should get a BEGIN at start."""
        from sltk.segmentation.postprocess import label_fixes

        # 5 consecutive INs after an OUT -- should trigger BEGIN insertion
        labels = torch.tensor([0, 1, 1, 1, 1, 1, 0])
        result = label_fixes(labels, min_segment_len=2)
        assert result[1].item() == 2  # BEGIN inserted

    def test_does_not_modify_input(self):
        from sltk.segmentation.postprocess import label_fixes

        labels = torch.tensor([0, 1, 1, 1, 0])
        original = labels.clone()
        label_fixes(labels)
        assert torch.equal(labels, original)

    def test_collapses_consecutive_begins(self):
        """Pass 2: consecutive BEGINs should collapse (first stays, rest become IN)."""
        from sltk.segmentation.postprocess import label_fixes

        labels = torch.tensor([2, 2, 2, 1, 0])
        result = label_fixes(labels)
        assert result[0].item() == 2
        assert result[1].item() == 1
        assert result[2].item() == 1

    def test_single_in(self):
        """A single IN should not get a BEGIN (below min_segment_len)."""
        from sltk.segmentation.postprocess import label_fixes

        labels = torch.tensor([0, 1, 0])
        result = label_fixes(labels, min_segment_len=2)
        # Single IN is below threshold, no BEGIN inserted
        assert 2 not in result.tolist()


class TestExtractSegments:
    """Test segment boundary extraction from BIO labels."""

    def test_no_segments(self):
        from sltk.segmentation.postprocess import extract_segments

        labels = torch.tensor([0, 0, 0, 0])
        result = extract_segments(labels)
        assert result == []

    def test_single_segment(self):
        from sltk.segmentation.postprocess import extract_segments

        labels = torch.tensor([0, 2, 1, 1, 0])
        result = extract_segments(labels)
        assert len(result) == 1
        assert result[0] == (1, 3)

    def test_two_segments_with_begin(self):
        from sltk.segmentation.postprocess import extract_segments

        labels = torch.tensor([2, 1, 1, 0, 2, 1, 0])
        result = extract_segments(labels)
        assert len(result) == 2
        assert result[0] == (0, 2)
        assert result[1] == (4, 5)

    def test_segment_at_end(self):
        from sltk.segmentation.postprocess import extract_segments

        labels = torch.tensor([0, 0, 1, 1, 1])
        result = extract_segments(labels)
        assert len(result) == 1
        assert result[0] == (2, 4)

    def test_all_in(self):
        from sltk.segmentation.postprocess import extract_segments

        labels = torch.tensor([1, 1, 1])
        result = extract_segments(labels)
        assert len(result) == 1
        assert result[0] == (0, 2)


class TestMergeBeginTags:
    """Test BEGIN -> IN merging."""

    def test_basic_merge(self):
        from sltk.segmentation.postprocess import merge_begin_tags

        labels = torch.tensor([0, 2, 1, 0, 2, 1])
        result = merge_begin_tags(labels)
        assert 2 not in result.tolist()
        assert result.tolist() == [0, 1, 1, 0, 1, 1]

    def test_does_not_modify_input(self):
        from sltk.segmentation.postprocess import merge_begin_tags

        labels = torch.tensor([2, 1, 0])
        original = labels.clone()
        merge_begin_tags(labels)
        assert torch.equal(labels, original)


# ===========================================================================
# H5 loader tests
# ===========================================================================


class TestMatrixToRotation6d:
    """Test rotation matrix to 6D conversion."""

    def test_identity_matrix(self):
        from sltk.segmentation.h5_loader import matrix_to_rotation_6d

        eye = np.eye(3, dtype=np.float32).reshape(1, 3, 3)
        result = matrix_to_rotation_6d(eye)
        assert result.shape == (1, 6)
        # First 6 values of identity's first two rows
        expected = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
        np.testing.assert_allclose(result[0], expected, atol=1e-6)

    def test_batch_shape(self):
        from sltk.segmentation.h5_loader import matrix_to_rotation_6d

        matrices = np.random.randn(5, 3, 3, 3).astype(np.float32)
        result = matrix_to_rotation_6d(matrices)
        assert result.shape == (5, 3, 6)

    def test_single_matrix(self):
        from sltk.segmentation.h5_loader import matrix_to_rotation_6d

        mat = np.random.randn(3, 3).astype(np.float32)
        result = matrix_to_rotation_6d(mat)
        assert result.shape == (6,)


# ===========================================================================
# Output tests
# ===========================================================================


class TestSaveAsJson:
    """Test JSON output format."""

    def test_saves_valid_json(self, tmp_path):
        from sltk.segmentation.output import save_as_json

        predictions = {
            "video1": torch.tensor([0, 2, 1, 1, 0, 0, 2, 1, 0]),
        }
        out = save_as_json(predictions, tmp_path / "out.json", fps=25.0)
        assert out.exists()

        with open(out) as f:
            data = json.load(f)

        assert "video1" in data
        v1 = data["video1"]
        assert v1["fps"] == 25.0
        assert v1["num_frames"] == 9
        assert len(v1["segments"]) == 2

    def test_timestamps_are_correct(self, tmp_path):
        from sltk.segmentation.output import save_as_json

        predictions = {"test": torch.tensor([0, 1, 1, 0])}
        out = save_as_json(predictions, tmp_path / "out.json", fps=10.0)

        with open(out) as f:
            data = json.load(f)

        seg = data["test"]["segments"][0]
        assert seg["start_frame"] == 1
        assert seg["end_frame"] == 2
        assert seg["start_sec"] == pytest.approx(0.1)
        assert seg["end_sec"] == pytest.approx(0.2)


class TestSaveAsElan:
    """Test ELAN EAF XML output."""

    def test_valid_xml(self, tmp_path):
        from sltk.segmentation.output import save_as_elan

        predictions = {"vid": torch.tensor([0, 2, 1, 1, 0])}
        out = save_as_elan(predictions, tmp_path / "out.eaf", fps=25.0)

        tree = ET.parse(out)
        root = tree.getroot()
        assert root.tag == "ANNOTATION_DOCUMENT"
        tiers = root.findall("TIER")
        assert len(tiers) == 1
        assert tiers[0].get("TIER_ID") == "vid_segmentation"

    def test_time_slots_exist(self, tmp_path):
        from sltk.segmentation.output import save_as_elan

        predictions = {"vid": torch.tensor([2, 1, 0, 2, 1])}
        out = save_as_elan(predictions, tmp_path / "out.eaf", fps=25.0)

        tree = ET.parse(out)
        root = tree.getroot()
        time_order = root.find("TIME_ORDER")
        slots = time_order.findall("TIME_SLOT")
        # 2 segments * 2 time slots each = 4
        assert len(slots) == 4


class TestSaveAsPt:
    """Test PyTorch output format."""

    def test_saves_and_loads(self, tmp_path):
        from sltk.segmentation.output import save_as_pt

        predictions = {"a": torch.tensor([0, 1, 2]), "b": torch.tensor([1, 0])}
        out = save_as_pt(predictions, tmp_path / "out.pt")
        assert out.exists()

        loaded = torch.load(out, weights_only=False)
        assert torch.equal(loaded["a"], predictions["a"])
        assert torch.equal(loaded["b"], predictions["b"])


class TestSavePredictions:
    """Test the save_predictions dispatcher."""

    def test_dispatches_json(self, tmp_path):
        from sltk.segmentation.output import OutputFormat, save_predictions

        predictions = {"x": torch.tensor([0, 1, 0])}
        out = save_predictions(predictions, tmp_path / "out.json", OutputFormat.JSON)
        assert out.exists()
        with open(out) as f:
            data = json.load(f)
        assert "x" in data

    def test_dispatches_elan(self, tmp_path):
        from sltk.segmentation.output import OutputFormat, save_predictions

        predictions = {"y": torch.tensor([2, 1, 0])}
        out = save_predictions(predictions, tmp_path / "out.eaf", OutputFormat.ELAN)
        assert out.exists()

    def test_dispatches_pt(self, tmp_path):
        from sltk.segmentation.output import OutputFormat, save_predictions

        predictions = {"z": torch.tensor([0])}
        out = save_predictions(predictions, tmp_path / "out.pt", OutputFormat.PT)
        assert out.exists()


# ===========================================================================
# Runner tests (mocked model)
# ===========================================================================


class TestInferenceRunnerConstruction:
    """Test InferenceRunner constructor and checkpoint resolution."""

    def test_raises_when_no_checkpoint(self):
        from sltk.segmentation.runner import InferenceRunner

        with patch("sltk.extraction.weights.get_weight_path", return_value=None):
            with pytest.raises(FileNotFoundError, match="Segmentation checkpoint"):
                InferenceRunner()

    def test_raises_for_missing_path(self, tmp_path):
        from sltk.segmentation.runner import InferenceRunner

        with pytest.raises(FileNotFoundError, match="Checkpoint not found"):
            InferenceRunner(checkpoint_path=tmp_path / "nonexistent.ckpt")


class TestGetRunner:
    """Test thread-safe singleton."""

    def test_singleton_returns_same_instance(self):
        import sltk.segmentation.runner as runner_mod

        mock_runner = MagicMock()
        original = runner_mod._runner
        try:
            runner_mod._runner = mock_runner
            result = runner_mod.get_runner()
            assert result is mock_runner
        finally:
            runner_mod._runner = original
