"""End-to-end integration tests -- full pipeline workflows combining export, quantize, validate, and benchmark."""

from __future__ import annotations

import os

import numpy as np
import pytest
import torch

from matrice_export import ExportPipeline
from matrice_export.adapters.onnx_input import OnnxInputAdapter
from matrice_export.quantization import benchmark_quantization, quantize_int8
from matrice_export.tracking import ExportTracker
from matrice_export.validators.openvino import OpenVinoValidator
from matrice_export.validators.tensorrt import TensorRTValidator

from .conftest import MockActionTracker

# ====================================================================== #
# Class 1: TestExportQuantizeBenchmarkE2E
# ====================================================================== #


@pytest.mark.e2e
class TestExportQuantizeBenchmarkE2E:
    """Full pipeline: export -> quantize -> benchmark."""

    def test_dummy_model_export_quantize_benchmark(
        self, dummy_model, sample_input, dummy_onnx_path, tmp_path
    ):
        """Export dummy model to ONNX, quantize to INT8, then benchmark both variants.

        Uses the dummy_onnx_path fixture (dynamo=False, opset 17) to avoid
        shape inference issues during quantization.
        """
        onnx_path = dummy_onnx_path

        # Step 1: Verify the ONNX file exists from the fixture
        assert os.path.isfile(onnx_path), f"ONNX model not found: {onnx_path}"

        # Step 2: Quantize to INT8
        int8_path = quantize_int8(
            onnx_path, output_path=str(tmp_path / "model_int8.onnx")
        )
        assert os.path.isfile(int8_path), f"INT8 model not found: {int8_path}"

        # Step 3: Benchmark quantization comparison
        sample_np = sample_input.numpy().astype(np.float32)
        report = benchmark_quantization(
            onnx_path,
            int8_path=int8_path,
            sample_input=sample_np,
        )

        # Assertions on report structure
        assert "original" in report, "Report missing 'original' key"
        assert "int8" in report, "Report missing 'int8' key"

        for variant_name in ("original", "int8"):
            variant = report[variant_name]
            # The dummy model is tiny (~500 bytes), so size_mb rounds to 0.00.
            # Verify size_mb is present and non-negative; also verify the file
            # actually exists on disk.
            assert variant["size_mb"] >= 0, (
                f"{variant_name} size_mb should be >= 0, got {variant['size_mb']}"
            )
            assert os.path.isfile(variant["path"]), (
                f"{variant_name} model file not found: {variant['path']}"
            )
            assert variant["latency"] is not None, (
                f"{variant_name} latency should not be None"
            )
            assert variant["latency"]["mean_ms"] > 0, (
                f"{variant_name} mean_ms should be > 0, got {variant['latency']['mean_ms']}"
            )

    @pytest.mark.real_model
    @pytest.mark.slow
    def test_resnet18_export_quantize_benchmark(
        self, resnet18_onnx_path, tmp_path
    ):
        """Export ResNet18 to ONNX, quantize to INT8, then benchmark.

        Uses the session-scoped resnet18_onnx_path fixture for speed.
        """
        # Step 1: Verify the ONNX file exists
        assert os.path.isfile(resnet18_onnx_path), (
            f"ResNet18 ONNX model not found: {resnet18_onnx_path}"
        )

        # Step 2: Quantize to INT8
        int8_path = quantize_int8(
            resnet18_onnx_path,
            output_path=str(tmp_path / "resnet18_int8.onnx"),
        )
        assert os.path.isfile(int8_path), f"INT8 model not found: {int8_path}"

        # Step 3: Benchmark
        sample_np = np.random.randn(1, 3, 224, 224).astype(np.float32)
        report = benchmark_quantization(
            resnet18_onnx_path,
            int8_path=int8_path,
            sample_input=sample_np,
        )

        # Assertions on report structure
        assert "original" in report
        assert "int8" in report

        # Original should always have valid latency
        original = report["original"]
        assert "path" in original
        assert original["size_mb"] > 0
        assert original["latency"] is not None
        assert "mean_ms" in original["latency"]
        assert original["latency"]["mean_ms"] > 0

        # INT8 may fail inference (e.g. ConvInteger not implemented in ORT)
        # — verify the benchmark captures this gracefully
        int8 = report["int8"]
        assert "path" in int8
        assert int8["size_mb"] > 0
        if isinstance(int8["latency"], dict) and "error" in int8["latency"]:
            # ORT doesn't support ConvInteger on this platform — graceful
            assert "error" in int8["latency"]
        else:
            assert int8["latency"] is not None
            assert "mean_ms" in int8["latency"]


# ====================================================================== #
# Class 2: TestExportValidateE2E
# ====================================================================== #


@pytest.mark.e2e
class TestExportValidateE2E:
    """Export with validation enabled -- verify that validators run inline."""

    def test_dummy_onnx_export_validate(self, dummy_model, sample_input, tmp_path):
        """Export dummy model to ONNX with validate=True, verify validation passes."""
        pipeline = ExportPipeline(
            dummy_model, sample_input, task="classification"
        )
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)

        assert results["onnx"]["status"] == "success"
        assert os.path.isfile(results["onnx"]["path"])

        validation = results["onnx"]["validation"]
        assert validation is not None, "Validation should not be None when validate=True"
        assert validation["shape_match"] is True
        assert validation["values_match"] is True

    def test_dummy_torchscript_export_validate(
        self, dummy_model, sample_input, tmp_path
    ):
        """Export dummy model to TorchScript with validate=True, verify validation passes."""
        pipeline = ExportPipeline(
            dummy_model, sample_input, task="classification"
        )
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)

        assert results["torchscript"]["status"] == "success"
        assert os.path.isfile(results["torchscript"]["path"])

        validation = results["torchscript"]["validation"]
        assert validation is not None, "Validation should not be None when validate=True"
        assert validation["shape_match"] is True
        assert validation["values_match"] is True

    def test_conv_model_multi_format_validate(
        self, conv_model, image_input, tmp_path
    ):
        """Export conv model to both ONNX and TorchScript with validate=True."""
        pipeline = ExportPipeline(
            conv_model, image_input, task="classification"
        )
        results = pipeline.export(
            ["onnx", "torchscript"], str(tmp_path), validate=True
        )

        for fmt in ("onnx", "torchscript"):
            assert results[fmt]["status"] == "success", (
                f"{fmt} export failed: {results[fmt].get('error')}"
            )
            assert os.path.isfile(results[fmt]["path"])

            validation = results[fmt]["validation"]
            assert validation is not None, (
                f"{fmt} validation should not be None"
            )
            assert validation["shape_match"] is True, (
                f"{fmt} shape_match is not True"
            )
            assert validation["values_match"] is True, (
                f"{fmt} values_match is not True"
            )

    @pytest.mark.real_model
    def test_resnet18_onnx_validate(
        self, resnet18_model, resnet18_input, tmp_path
    ):
        """Export ResNet18 to ONNX with validate=True, verify shape_match=True."""
        pipeline = ExportPipeline(
            resnet18_model, resnet18_input, task="classification"
        )
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)

        assert results["onnx"]["status"] == "success"
        assert os.path.isfile(results["onnx"]["path"])

        validation = results["onnx"]["validation"]
        assert validation is not None
        assert validation["shape_match"] is True


# ====================================================================== #
# Class 3: TestExportWithTrackingE2E
# ====================================================================== #


@pytest.mark.e2e
class TestExportWithTrackingE2E:
    """Test pipeline with ExportTracker -- verify all step codes are emitted."""

    def test_tracked_export_sends_all_step_codes(
        self, dummy_model, sample_input, tmp_path
    ):
        """Run a tracked export and verify all expected step codes are sent.

        Expected step codes in order:
          MDL_EXP_ACK  -- pipeline acknowledged
          MDL_EXP_STR  -- pipeline started
          MDL_EXPT_ACK -- per-format export acknowledged
          MDL_EXPT     -- per-format export result
          MDL_EXP_CMPL -- pipeline completed
        """
        mock_at = MockActionTracker()
        tracker = ExportTracker(action_tracker=mock_at)
        pipeline = ExportPipeline(
            dummy_model, sample_input, task="classification", tracker=tracker
        )

        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        assert results["onnx"]["status"] == "success"

        # Extract all step codes in order
        step_codes = [call[0] for call in mock_at.calls]

        expected_codes = [
            "MDL_EXP_ACK",
            "MDL_EXP_STR",
            "MDL_EXPT_ACK",
            "MDL_EXPT",
            "MDL_EXP_CMPL",
        ]
        for code in expected_codes:
            assert code in step_codes, (
                f"Missing step code {code!r} in tracker calls. "
                f"Got: {step_codes}"
            )

    def test_tracked_export_success_status(
        self, dummy_model, sample_input, tmp_path
    ):
        """Verify that MDL_EXPT has SUCCESS status and MDL_EXP_CMPL has SUCCESS."""
        mock_at = MockActionTracker()
        tracker = ExportTracker(action_tracker=mock_at)
        pipeline = ExportPipeline(
            dummy_model, sample_input, task="classification", tracker=tracker
        )

        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        assert results["onnx"]["status"] == "success"

        # Check MDL_EXPT has SUCCESS status
        expt_calls = [c for c in mock_at.calls if c[0] == "MDL_EXPT"]
        assert len(expt_calls) >= 1, "No MDL_EXPT call found"
        assert expt_calls[0][1] == "SUCCESS", (
            f"MDL_EXPT status should be SUCCESS, got {expt_calls[0][1]!r}"
        )

        # Check MDL_EXP_CMPL has SUCCESS status
        cmpl_calls = [c for c in mock_at.calls if c[0] == "MDL_EXP_CMPL"]
        assert len(cmpl_calls) >= 1, "No MDL_EXP_CMPL call found"
        assert cmpl_calls[0][1] == "SUCCESS", (
            f"MDL_EXP_CMPL status should be SUCCESS, got {cmpl_calls[0][1]!r}"
        )


# ====================================================================== #
# Class 4: TestTensorRTPipelineE2E
# ====================================================================== #


# Conditional imports for GPU/TensorRT tests
try:
    import tensorrt as trt  # noqa: F401

    _HAS_TRT = True
    _TRT_MAJOR = int(trt.__version__.split(".")[0])
except (ImportError, Exception):
    _HAS_TRT = False
    _TRT_MAJOR = 0

try:
    import pycuda.driver  # noqa: F401

    _HAS_PYCUDA = True
except ImportError:
    _HAS_PYCUDA = False

_SKIP_TRT = not (_HAS_TRT and _HAS_PYCUDA and torch.cuda.is_available())
_TRT_REASON = "Requires CUDA GPU + tensorrt + pycuda"
_TRT10_XFAIL = _TRT_MAJOR >= 10


@pytest.mark.e2e
@pytest.mark.gpu
@pytest.mark.tensorrt
@pytest.mark.skipif(_SKIP_TRT, reason=_TRT_REASON)
class TestTensorRTPipelineE2E:
    """Full GPU pipeline: ONNX -> TensorRT engine -> validate."""

    @pytest.mark.timeout(180)
    @pytest.mark.xfail(
        _TRT10_XFAIL,
        reason="TRT 10+ removed num_bindings API; validator needs update",
        strict=False,
    )
    def test_dummy_onnx_to_tensorrt_to_validate(
        self, dummy_onnx_path, dummy_model, sample_input, tmp_path
    ):
        """Convert dummy ONNX to TensorRT engine and validate against PyTorch baseline.

        Steps:
          1. Use dummy_onnx_path fixture (Linear(10,5) exported with dynamo=False)
          2. Convert to TensorRT engine via OnnxInputAdapter
          3. Validate with TensorRTValidator against PyTorch baseline
          4. Assert shape_match=True, latency_ms > 0
        """
        # Generate baseline from PyTorch model
        dummy_model.eval()
        np_input = sample_input.numpy().astype(np.float32)
        with torch.no_grad():
            baseline = dummy_model(sample_input).numpy()

        # Convert ONNX to TensorRT engine
        adapter = OnnxInputAdapter()
        engine_path = adapter.to_tensorrt(
            dummy_onnx_path, str(tmp_path), half=False, workspace=1
        )
        assert os.path.isfile(engine_path), (
            f"TensorRT engine file not found: {engine_path}"
        )

        # Validate engine
        result = TensorRTValidator().validate(
            engine_path, np_input, baseline=baseline
        )

        # Must not be skipped or error
        assert "status" not in result or result.get("status") not in (
            "skipped",
            "error",
        ), f"Unexpected status in result: {result}"

        assert result["shape_match"] is True, (
            f"shape_match should be True, got {result['shape_match']}"
        )
        assert result["latency_ms"] > 0, (
            f"latency_ms should be > 0, got {result['latency_ms']}"
        )


# ====================================================================== #
# Class 5: TestOpenVINOPipelineE2E
# ====================================================================== #


# Conditional import for OpenVINO
try:
    from openvino.runtime import Core as _OVCore  # noqa: F401

    _HAS_OV = True
except ImportError:
    _HAS_OV = False

_SKIP_OV = not _HAS_OV
_OV_REASON = "Requires openvino>=2023.0"


@pytest.mark.e2e
@pytest.mark.openvino
@pytest.mark.skipif(_SKIP_OV, reason=_OV_REASON)
class TestOpenVINOPipelineE2E:
    """Full OpenVINO pipeline: ONNX -> OpenVINO IR -> validate."""

    def test_dummy_onnx_to_openvino_to_validate(
        self, dummy_onnx_path, dummy_model, sample_input, tmp_path
    ):
        """Convert dummy ONNX to OpenVINO IR and validate against PyTorch baseline.

        Steps:
          1. Use dummy_onnx_path fixture (Linear(10,5) exported with dynamo=False)
          2. Convert to OpenVINO IR via OnnxInputAdapter
          3. Validate with OpenVinoValidator against PyTorch baseline
          4. Assert shape_match=True, latency_ms > 0
        """
        # Generate baseline from PyTorch model
        dummy_model.eval()
        np_input = sample_input.numpy().astype(np.float32)
        with torch.no_grad():
            baseline = dummy_model(sample_input).numpy()

        # Convert ONNX to OpenVINO IR
        adapter = OnnxInputAdapter()
        openvino_dir = adapter.to_openvino(
            dummy_onnx_path, str(tmp_path / "ov_out")
        )
        assert os.path.isdir(openvino_dir), (
            f"OpenVINO model directory not found: {openvino_dir}"
        )

        # Validate OpenVINO model
        result = OpenVinoValidator().validate(
            openvino_dir, np_input, baseline=baseline
        )

        # Must not be skipped or error
        assert "status" not in result or result.get("status") not in (
            "skipped",
            "error",
        ), f"Unexpected status in result: {result}"

        assert result["shape_match"] is True, (
            f"shape_match should be True, got {result['shape_match']}"
        )
        assert result["latency_ms"] > 0, (
            f"latency_ms should be > 0, got {result['latency_ms']}"
        )
