"""Concrete validator class tests — OnnxValidator and TorchScriptValidator.

Unlike test_validators.py which uses inline validator functions, these tests
exercise the actual validator classes from matrice_export.validators.
"""

from __future__ import annotations

import pytest
import torch

from matrice_export import ExportPipeline
from matrice_export.validators.onnx_val import OnnxValidator
from matrice_export.validators.torchscript import TorchScriptValidator

# ------------------------------------------------------------------ #
# Fixtures
# ------------------------------------------------------------------ #


@pytest.fixture
def exported_onnx(dummy_model, sample_input, tmp_path):
    """Export dummy model to ONNX and return (path, baseline_output)."""
    dummy_model.eval()
    with torch.no_grad():
        baseline = dummy_model(sample_input).numpy()
    pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
    results = pipeline.export(["onnx"], str(tmp_path), validate=False)
    return results["onnx"]["path"], baseline


@pytest.fixture
def exported_torchscript(dummy_model, sample_input, tmp_path):
    """Export dummy model to TorchScript and return (path, baseline_output)."""
    dummy_model.eval()
    with torch.no_grad():
        baseline = dummy_model(sample_input).numpy()
    pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
    results = pipeline.export(["torchscript"], str(tmp_path), validate=False)
    return results["torchscript"]["path"], baseline


# ------------------------------------------------------------------ #
# 1. OnnxValidator concrete tests
# ------------------------------------------------------------------ #


class TestOnnxValidatorConcrete:
    def test_validate_returns_all_keys(self, exported_onnx, sample_input):
        """OnnxValidator.validate returns dict with all expected keys."""
        path, baseline = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=baseline)

        assert "shape_match" in result
        assert "values_match" in result
        assert "max_diff" in result
        assert "latency_ms" in result
        assert "output_shape" in result

    def test_shape_match_true(self, exported_onnx, sample_input):
        """For a correctly exported ONNX model, shape_match is True."""
        path, baseline = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=baseline)
        assert result["shape_match"] is True

    def test_values_match_true(self, exported_onnx, sample_input):
        """For a simple model, ONNX values closely match PyTorch baseline."""
        path, baseline = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=baseline)
        assert result["values_match"] is True

    def test_latency_positive(self, exported_onnx, sample_input):
        """Latency measurement is a positive number."""
        path, baseline = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=baseline)
        assert result["latency_ms"] > 0

    def test_max_diff_small(self, exported_onnx, sample_input):
        """max_diff should be very small for a simple model."""
        path, baseline = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=baseline)
        assert result["max_diff"] < 1e-4


# ------------------------------------------------------------------ #
# 2. TorchScriptValidator concrete tests
# ------------------------------------------------------------------ #


class TestTorchScriptValidatorConcrete:
    def test_validate_returns_all_keys(self, exported_torchscript, sample_input):
        """TorchScriptValidator.validate returns dict with all expected keys."""
        path, baseline = exported_torchscript
        np_input = sample_input.numpy()
        result = TorchScriptValidator().validate(path, np_input, baseline=baseline)

        assert "shape_match" in result
        assert "values_match" in result
        assert "max_diff" in result
        assert "latency_ms" in result
        assert "output_shape" in result

    def test_shape_match_true(self, exported_torchscript, sample_input):
        """For a correctly exported TorchScript model, shape_match is True."""
        path, baseline = exported_torchscript
        np_input = sample_input.numpy()
        result = TorchScriptValidator().validate(path, np_input, baseline=baseline)
        assert result["shape_match"] is True

    def test_values_match_true(self, exported_torchscript, sample_input):
        """For a simple model, TorchScript values closely match baseline."""
        path, baseline = exported_torchscript
        np_input = sample_input.numpy()
        result = TorchScriptValidator().validate(path, np_input, baseline=baseline)
        assert result["values_match"] is True

    def test_latency_positive(self, exported_torchscript, sample_input):
        """Latency measurement is a positive number."""
        path, baseline = exported_torchscript
        np_input = sample_input.numpy()
        result = TorchScriptValidator().validate(path, np_input, baseline=baseline)
        assert result["latency_ms"] > 0


# ------------------------------------------------------------------ #
# 3. Validate without baseline
# ------------------------------------------------------------------ #


class TestValidatorNoBaseline:
    def test_onnx_no_baseline_shape_none(self, exported_onnx, sample_input):
        """Without baseline, shape_match and values_match are None."""
        path, _ = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=None)
        assert result["shape_match"] is None
        assert result["values_match"] is None
        assert result["max_diff"] is None

    def test_torchscript_no_baseline_shape_none(self, exported_torchscript, sample_input):
        """Without baseline, shape_match and values_match are None."""
        path, _ = exported_torchscript
        np_input = sample_input.numpy()
        result = TorchScriptValidator().validate(path, np_input, baseline=None)
        assert result["shape_match"] is None
        assert result["values_match"] is None

    def test_onnx_no_baseline_still_has_output_shape(self, exported_onnx, sample_input):
        """Without baseline, output_shape is still present."""
        path, _ = exported_onnx
        np_input = sample_input.numpy()
        result = OnnxValidator().validate(path, np_input, baseline=None)
        assert "output_shape" in result
        assert result["output_shape"] == (1, 5)


# ------------------------------------------------------------------ #
# 4. Bad model path
# ------------------------------------------------------------------ #


class TestValidatorBadPath:
    def test_onnx_bad_path_returns_error(self, sample_input):
        """Non-existent model path returns an error dict."""
        np_input = sample_input.numpy()
        result = OnnxValidator().validate("/tmp/nonexistent_model_xyz.onnx", np_input)
        assert result["status"] == "error"

    def test_torchscript_bad_path_returns_error(self, sample_input):
        """Non-existent model path returns an error dict."""
        np_input = sample_input.numpy()
        result = TorchScriptValidator().validate("/tmp/nonexistent_model_xyz.torchscript", np_input)
        assert result["status"] == "error"
