"""Validator tests — structure, correctness, and latency for ONNX and TorchScript validators.

Since the validator module (matrice_export.validators) does not yet have concrete
implementations registered, these tests define inline validator functions that
exercise the validator registration and invocation path of ExportPipeline.
They also verify the expected dict structure and correctness guarantees that
any production validator must satisfy.
"""

from __future__ import annotations

import time

import numpy as np
import onnxruntime as ort
import pytest
import torch

from matrice_export import ExportPipeline

# ------------------------------------------------------------------ #
# Helper: inline ONNX validator
# ------------------------------------------------------------------ #

def _onnx_validator(model_path: str, sample_input: np.ndarray, baseline: np.ndarray | None, **kw) -> dict:
    """A minimal ONNX validator that checks shape, values, and measures latency."""
    sess = ort.InferenceSession(model_path)
    input_name = sess.get_inputs()[0].name

    t0 = time.perf_counter()
    onnx_output = sess.run(None, {input_name: sample_input})[0]
    latency_ms = (time.perf_counter() - t0) * 1000

    shape_match = True
    values_match = True

    if baseline is not None:
        shape_match = onnx_output.shape == baseline.shape
        values_match = bool(np.allclose(onnx_output, baseline, rtol=1e-4, atol=1e-4))

    return {
        "status": "pass" if (shape_match and values_match) else "fail",
        "shape_match": shape_match,
        "values_match": values_match,
        "latency_ms": latency_ms,
        "output_shape": list(onnx_output.shape),
    }


# ------------------------------------------------------------------ #
# Helper: inline TorchScript validator
# ------------------------------------------------------------------ #

def _torchscript_validator(model_path: str, sample_input: np.ndarray, baseline: np.ndarray | None, **kw) -> dict:
    """A minimal TorchScript validator that checks shape, values, and measures latency."""
    ts_model = torch.jit.load(model_path)
    ts_model.eval()

    tensor_input = torch.from_numpy(sample_input)

    t0 = time.perf_counter()
    with torch.no_grad():
        ts_output = ts_model(tensor_input).numpy()
    latency_ms = (time.perf_counter() - t0) * 1000

    shape_match = True
    values_match = True

    if baseline is not None:
        shape_match = ts_output.shape == baseline.shape
        values_match = bool(np.allclose(ts_output, baseline, rtol=1e-4, atol=1e-4))

    return {
        "status": "pass" if (shape_match and values_match) else "fail",
        "shape_match": shape_match,
        "values_match": values_match,
        "latency_ms": latency_ms,
        "output_shape": list(ts_output.shape),
    }


# ------------------------------------------------------------------ #
# Fixture: register validators for test scope, restore afterwards
# ------------------------------------------------------------------ #

@pytest.fixture(autouse=True)
def _register_validators():
    """Temporarily register inline validators and restore original state after each test."""
    saved = ExportPipeline.VALIDATORS.copy()
    ExportPipeline.register_validator("onnx", _onnx_validator)
    ExportPipeline.register_validator("torchscript", _torchscript_validator)
    yield
    ExportPipeline.VALIDATORS.clear()
    ExportPipeline.VALIDATORS.update(saved)


# ------------------------------------------------------------------ #
# 1. ONNX validator returns correct dict structure
# ------------------------------------------------------------------ #

class TestOnnxValidatorStructure:
    def test_onnx_validation_dict_keys(self, dummy_model, sample_input, tmp_path):
        """The ONNX validator returns a dict with the expected keys."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)

        validation = results["onnx"]["validation"]
        assert validation is not None
        assert "shape_match" in validation
        assert "values_match" in validation
        assert "latency_ms" in validation
        assert "output_shape" in validation


# ------------------------------------------------------------------ #
# 2. TorchScript validator returns correct dict structure
# ------------------------------------------------------------------ #

class TestTorchScriptValidatorStructure:
    def test_torchscript_validation_dict_keys(self, dummy_model, sample_input, tmp_path):
        """The TorchScript validator returns a dict with the expected keys."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)

        validation = results["torchscript"]["validation"]
        assert validation is not None
        assert "status" in validation
        assert "shape_match" in validation
        assert "values_match" in validation
        assert "latency_ms" in validation
        assert "output_shape" in validation


# ------------------------------------------------------------------ #
# 3. Validator shape_match is True for correctly exported model
# ------------------------------------------------------------------ #

class TestValidatorShapeMatch:
    def test_onnx_shape_match_true(self, dummy_model, sample_input, tmp_path):
        """For a correctly exported ONNX model, shape_match is True."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)
        assert results["onnx"]["validation"]["shape_match"] is True

    def test_torchscript_shape_match_true(self, dummy_model, sample_input, tmp_path):
        """For a correctly exported TorchScript model, shape_match is True."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)
        assert results["torchscript"]["validation"]["shape_match"] is True


# ------------------------------------------------------------------ #
# 4. Validator values_match is True (small model, tight tolerance)
# ------------------------------------------------------------------ #

class TestValidatorValuesMatch:
    def test_onnx_values_match_true(self, dummy_model, sample_input, tmp_path):
        """For a simple model, ONNX values closely match PyTorch baseline."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)
        assert results["onnx"]["validation"]["values_match"] is True

    def test_torchscript_values_match_true(self, dummy_model, sample_input, tmp_path):
        """For a simple model, TorchScript values closely match PyTorch baseline."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)
        assert results["torchscript"]["validation"]["values_match"] is True


# ------------------------------------------------------------------ #
# 5. Validator latency_ms is positive
# ------------------------------------------------------------------ #

class TestValidatorLatency:
    def test_onnx_latency_positive(self, dummy_model, sample_input, tmp_path):
        """ONNX inference latency reported by the validator is a positive number."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)
        assert results["onnx"]["validation"]["latency_ms"] > 0

    def test_torchscript_latency_positive(self, dummy_model, sample_input, tmp_path):
        """TorchScript inference latency reported by the validator is a positive number."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)
        assert results["torchscript"]["validation"]["latency_ms"] > 0


# ------------------------------------------------------------------ #
# Additional: overall status is "pass"
# ------------------------------------------------------------------ #

class TestValidatorOverallStatus:
    def test_onnx_status_pass(self, dummy_model, sample_input, tmp_path):
        """When shape and values match, the ONNX validator reports status='pass'."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)
        assert results["onnx"]["validation"]["status"] == "pass"

    def test_torchscript_status_pass(self, dummy_model, sample_input, tmp_path):
        """When shape and values match, the TorchScript validator reports status='pass'."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)
        assert results["torchscript"]["validation"]["status"] == "pass"
