"""OpenVINO end-to-end tests -- ONNX to OpenVINO conversion and validation."""

from __future__ import annotations

import os
from unittest.mock import patch

import pytest
import torch

from matrice_export.adapters.onnx_input import OnnxInputAdapter
from matrice_export.validators.openvino import OpenVinoValidator

# ------------------------------------------------------------------ #
# Local fixtures
# ------------------------------------------------------------------ #


@pytest.fixture
def adapter():
    """Return an OnnxInputAdapter instance."""
    return OnnxInputAdapter()


@pytest.fixture
def openvino_model(adapter, dummy_onnx_path, tmp_path):
    """Convert the dummy ONNX model to OpenVINO IR and return (openvino_dir, xml_path).

    Uses the dummy_onnx_path fixture (small Linear+ReLU model exported
    with dynamo=False, opset 17) and converts it to OpenVINO IR format.
    """
    ov_dir = adapter.to_openvino(dummy_onnx_path, str(tmp_path))
    # The stem of the ONNX file is "dummy", so the xml lives under
    # <tmp_path>/dummy_openvino_model/dummy.xml
    xml_files = [f for f in os.listdir(ov_dir) if f.endswith(".xml")]
    assert len(xml_files) == 1, f"Expected exactly 1 .xml file, found: {xml_files}"
    xml_path = os.path.join(ov_dir, xml_files[0])
    return ov_dir, xml_path


@pytest.fixture
def dummy_baseline(dummy_model, sample_input):
    """Compute the PyTorch baseline output for the dummy model."""
    dummy_model.eval()
    with torch.no_grad():
        output = dummy_model(sample_input).numpy()
    return output


# ------------------------------------------------------------------ #
# 1. Conversion tests
# ------------------------------------------------------------------ #


@pytest.mark.openvino
class TestOpenVinoConversion:
    """Test OnnxInputAdapter.to_openvino() conversion from ONNX to OpenVINO IR."""

    def test_convert_dummy_onnx_to_openvino(self, adapter, dummy_onnx_path, tmp_path):
        """Converting a valid ONNX model produces an output directory containing IR files."""
        ov_dir = adapter.to_openvino(dummy_onnx_path, str(tmp_path))

        assert os.path.isdir(ov_dir)
        contents = os.listdir(ov_dir)
        xml_files = [f for f in contents if f.endswith(".xml")]
        bin_files = [f for f in contents if f.endswith(".bin")]
        assert len(xml_files) >= 1, f"No .xml file found in {ov_dir}: {contents}"
        assert len(bin_files) >= 1, f"No .bin file found in {ov_dir}: {contents}"

    def test_convert_produces_xml_and_bin(self, openvino_model):
        """Both .xml and .bin files exist with non-zero size."""
        ov_dir, xml_path = openvino_model

        bin_path = xml_path.replace(".xml", ".bin")

        assert os.path.isfile(xml_path), f".xml not found: {xml_path}"
        assert os.path.isfile(bin_path), f".bin not found: {bin_path}"
        assert os.path.getsize(xml_path) > 0, ".xml file is empty"
        assert os.path.getsize(bin_path) > 0, ".bin file is empty"

    def test_convert_with_fp16_compression(self, adapter, dummy_onnx_path, tmp_path):
        """Converting with half=True produces valid IR files (FP16 compressed)."""
        ov_dir = adapter.to_openvino(dummy_onnx_path, str(tmp_path), half=True)

        assert os.path.isdir(ov_dir)
        contents = os.listdir(ov_dir)
        xml_files = [f for f in contents if f.endswith(".xml")]
        bin_files = [f for f in contents if f.endswith(".bin")]
        assert len(xml_files) >= 1
        assert len(bin_files) >= 1

        # FP16 bin should still be non-empty
        bin_path = os.path.join(ov_dir, bin_files[0])
        assert os.path.getsize(bin_path) > 0

    def test_convert_nonexistent_onnx_raises(self, adapter, tmp_path):
        """Passing a non-existent ONNX path raises FileNotFoundError.

        Note: The FileNotFoundError is raised by the adapter.export()
        method which wraps to_openvino(). Calling to_openvino() directly
        may raise a different error from openvino. We test via export()
        to exercise the documented contract.
        """
        bad_path = os.path.join(str(tmp_path), "nonexistent.onnx")
        with pytest.raises(FileNotFoundError):
            adapter.export(bad_path, "openvino", str(tmp_path))


# ------------------------------------------------------------------ #
# 2. Validator tests
# ------------------------------------------------------------------ #


@pytest.mark.openvino
class TestOpenVinoValidator:
    """Test OpenVinoValidator with properly converted models."""

    def test_validate_returns_all_keys(self, openvino_model, sample_input, dummy_baseline):
        """Validation result contains all expected keys."""
        ov_dir, _ = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(ov_dir, np_input, baseline=dummy_baseline)

        expected_keys = {"shape_match", "values_match", "max_diff", "latency_ms", "output_shape"}
        for key in expected_keys:
            assert key in result, f"Missing key {key!r} in result: {result}"

    def test_validate_shape_match_true(self, openvino_model, sample_input, dummy_baseline):
        """Output shape from OpenVINO inference matches the PyTorch baseline."""
        ov_dir, _ = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(ov_dir, np_input, baseline=dummy_baseline)

        assert result["shape_match"] is True

    def test_validate_values_match_true(self, openvino_model, sample_input, dummy_baseline):
        """OpenVINO inference values match PyTorch baseline within atol=1e-2."""
        ov_dir, _ = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(
            ov_dir, np_input, baseline=dummy_baseline, atol=1e-2, rtol=1e-2,
        )

        assert result["values_match"] is True

    def test_validate_latency_positive(self, openvino_model, sample_input, dummy_baseline):
        """Measured latency is a positive number (in milliseconds)."""
        ov_dir, _ = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(ov_dir, np_input, baseline=dummy_baseline)

        assert result["latency_ms"] > 0

    def test_validate_with_directory_path(self, openvino_model, sample_input, dummy_baseline):
        """Passing the directory (not .xml) as model_path works correctly.

        The validator should auto-discover the .xml inside the directory.
        """
        ov_dir, _ = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(ov_dir, np_input, baseline=dummy_baseline)

        assert "status" not in result or result.get("status") != "error"
        assert result["shape_match"] is True

    def test_validate_with_xml_path(self, openvino_model, sample_input, dummy_baseline):
        """Passing the .xml file directly as model_path also works."""
        _, xml_path = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(xml_path, np_input, baseline=dummy_baseline)

        assert "status" not in result or result.get("status") != "error"
        assert result["shape_match"] is True

    def test_validate_without_baseline(self, openvino_model, sample_input):
        """Without a baseline, shape_match and values_match are None but output_shape is present."""
        ov_dir, _ = openvino_model
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(ov_dir, np_input, baseline=None)

        assert result["shape_match"] is None
        assert result["values_match"] is None
        assert result["max_diff"] is None
        assert "output_shape" in result
        assert isinstance(result["output_shape"], tuple)
        # The dummy model is Linear(10, 5) -> ReLU, so output shape is (1, 5)
        assert result["output_shape"] == (1, 5)


# ------------------------------------------------------------------ #
# 3. Skipped-when-not-installed tests
# ------------------------------------------------------------------ #


class TestOpenVinoValidatorSkips:
    """Test graceful skip behavior when OpenVINO is not installed."""

    def test_validate_without_openvino_returns_skipped(self, sample_input, tmp_path):
        """When openvino import fails, validator returns a skipped status dict."""
        np_input = sample_input.numpy()

        # Simulate openvino not being installed by making the import fail
        import builtins

        real_import = builtins.__import__

        def mock_import(name, *args, **kwargs):
            if name == "openvino.runtime" or name.startswith("openvino"):
                raise ImportError("Mocked: openvino not installed")
            return real_import(name, *args, **kwargs)

        with patch("builtins.__import__", side_effect=mock_import):
            result = OpenVinoValidator().validate(
                str(tmp_path), np_input, baseline=None,
            )

        assert result["status"] == "skipped"
        assert "openvino" in result["reason"].lower()


# ------------------------------------------------------------------ #
# 4. Validator error-handling tests
# ------------------------------------------------------------------ #


@pytest.mark.openvino
class TestOpenVinoValidatorErrors:
    """Test error handling in OpenVinoValidator."""

    def test_validate_bad_path_returns_error(self, sample_input):
        """A non-existent model path returns status=error."""
        np_input = sample_input.numpy()

        result = OpenVinoValidator().validate(
            "/tmp/nonexistent_openvino_model_xyz.xml", np_input,
        )

        assert result["status"] == "error"
        assert "error" in result

    def test_validate_empty_directory_returns_error(self, sample_input, tmp_path):
        """An empty directory (no .xml inside) returns status=error with descriptive message."""
        np_input = sample_input.numpy()

        # Create an empty subdirectory to pass as the model path
        empty_dir = os.path.join(str(tmp_path), "empty_openvino_model")
        os.makedirs(empty_dir, exist_ok=True)

        result = OpenVinoValidator().validate(empty_dir, np_input)

        assert result["status"] == "error"
        assert "error" in result
        assert ".xml" in result["error"].lower() or "xml" in result["error"].lower()
