"""
GPU integration tests for the segmentation v2 pipeline.

These tests require:
- A CUDA-capable GPU
- The segmentor v2 checkpoint at sltk/weights/segmentor/segmentor_v2.ckpt
- Example WiLoR H5 file at tests/fixtures/example_wilor.h5

Run with:
    pytest tests/test_segmentation_integration.py -v -m gpu
    pytest tests/test_segmentation_integration.py -v -m "gpu and not slow"
"""

from __future__ import annotations

from pathlib import Path

import pytest

torch = pytest.importorskip("torch")

CHECKPOINT = Path("/home/ef0036/Projects/VLT/sltk/weights/segmentor/segmentor_v2.ckpt")
EXAMPLE_H5 = Path("/home/ef0036/Projects/VLT/tests/fixtures/example_wilor.h5")

pytestmark = [
    pytest.mark.gpu,
    pytest.mark.slow,
    pytest.mark.skipif(
        not torch.cuda.is_available(),
        reason="GPU required",
    ),
    pytest.mark.skipif(
        not CHECKPOINT.exists(),
        reason=f"Checkpoint not found: {CHECKPOINT}",
    ),
    pytest.mark.skipif(
        not EXAMPLE_H5.exists(),
        reason=f"Example H5 not found: {EXAMPLE_H5}",
    ),
]


# ---------------------------------------------------------------------------
# Module-scoped runner fixture (expensive -- load once)
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def runner():
    """Load the real segmentation runner on GPU (once per module)."""
    from sltk.segmentation.runner import InferenceRunner

    return InferenceRunner(checkpoint_path=CHECKPOINT)


# ---------------------------------------------------------------------------
# H5 loading tests
# ---------------------------------------------------------------------------


class TestH5Loading:
    """Test loading WiLoR H5 features."""

    def test_h5_to_features_shape(self):
        from sltk.segmentation.h5_loader import h5_to_features

        features = h5_to_features(EXAMPLE_H5)
        assert features.dim() == 2
        assert features.shape[1] == 192
        assert features.shape[0] > 0

    def test_h5_to_features_dtype(self):
        from sltk.segmentation.h5_loader import h5_to_features

        features = h5_to_features(EXAMPLE_H5)
        assert features.dtype == torch.float32


# ---------------------------------------------------------------------------
# Full pipeline tests
# ---------------------------------------------------------------------------


class TestFullPipeline:
    """End-to-end: H5 -> features -> BIO labels -> segments."""

    def test_predict_returns_bio_labels(self, runner):
        from sltk.segmentation.h5_loader import h5_to_features

        features = h5_to_features(EXAMPLE_H5)
        labels = runner.predict(features, include_begin=True)

        assert labels.dim() == 1
        assert labels.shape[0] == features.shape[0]
        unique_vals = set(labels.unique().tolist())
        assert unique_vals.issubset({0, 1, 2})

    def test_predict_without_begin(self, runner):
        from sltk.segmentation.h5_loader import h5_to_features

        features = h5_to_features(EXAMPLE_H5)
        labels = runner.predict(features, include_begin=False)
        assert 2 not in labels.tolist()

    def test_extract_segments_from_predictions(self, runner):
        from sltk.segmentation.h5_loader import h5_to_features
        from sltk.segmentation.postprocess import extract_segments

        features = h5_to_features(EXAMPLE_H5)
        labels = runner.predict(features, include_begin=True)
        segments = extract_segments(labels)

        assert isinstance(segments, list)
        # Should find at least one segment in the example
        assert len(segments) > 0
        for start, end in segments:
            assert 0 <= start <= end < labels.shape[0]

    def test_run_h5_single_file(self, runner):
        predictions = runner.run_h5(EXAMPLE_H5)
        assert len(predictions) == 1
        name = EXAMPLE_H5.stem
        assert name in predictions
        assert predictions[name].dim() == 1

    def test_run_h5_directory(self, runner):
        predictions = runner.run_h5(EXAMPLE_H5.parent)
        assert len(predictions) >= 1


# ---------------------------------------------------------------------------
# Output format tests
# ---------------------------------------------------------------------------


class TestOutputFormats:
    """Test saving predictions in different formats."""

    def test_save_json(self, runner, tmp_path):
        import json

        from sltk.segmentation.output import OutputFormat, save_predictions

        predictions = runner.run_h5(EXAMPLE_H5)
        out = save_predictions(predictions, tmp_path / "out.json", OutputFormat.JSON)
        assert out.exists()

        with open(out) as f:
            data = json.load(f)
        assert EXAMPLE_H5.stem in data

    def test_save_elan(self, runner, tmp_path):
        import xml.etree.ElementTree as ET

        from sltk.segmentation.output import OutputFormat, save_predictions

        predictions = runner.run_h5(EXAMPLE_H5)
        out = save_predictions(predictions, tmp_path / "out.eaf", OutputFormat.ELAN)
        assert out.exists()

        tree = ET.parse(out)
        root = tree.getroot()
        assert root.tag == "ANNOTATION_DOCUMENT"

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

        predictions = runner.run_h5(EXAMPLE_H5)
        out = save_predictions(predictions, tmp_path / "out.pt", OutputFormat.PT)
        assert out.exists()

        loaded = torch.load(out, weights_only=False)
        for name, tensor in loaded.items():
            assert tensor.dim() == 1
