"""Real model export tests -- ResNet18 and YOLOv8n with actual weights."""

from __future__ import annotations

import os

import pytest
import torch

from matrice_export import ExportPipeline
from matrice_export.adapters.ultralytics import UltralyticsAdapter

# ------------------------------------------------------------------ #
# 1. ResNet18 exports via ExportPipeline
# ------------------------------------------------------------------ #


@pytest.mark.real_model
class TestResNet18Export:
    """Export a pretrained ResNet18 through the ExportPipeline."""

    def test_export_to_onnx(self, resnet18_model, resnet18_input, tmp_path):
        """ExportPipeline exports ResNet18 to ONNX successfully."""
        pipeline = ExportPipeline(
            resnet18_model, resnet18_input, task="classification"
        )
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        assert "onnx" in results
        assert results["onnx"]["status"] == "success"
        assert results["onnx"]["path"] is not None
        assert os.path.isfile(results["onnx"]["path"])

    def test_export_to_torchscript(self, resnet18_model, resnet18_input, tmp_path):
        """ExportPipeline exports ResNet18 to TorchScript successfully."""
        pipeline = ExportPipeline(
            resnet18_model, resnet18_input, task="classification"
        )
        results = pipeline.export(["torchscript"], str(tmp_path), validate=False)

        assert "torchscript" in results
        assert results["torchscript"]["status"] == "success"
        assert results["torchscript"]["path"] is not None
        assert os.path.isfile(results["torchscript"]["path"])
        # Verify the TorchScript file is loadable
        loaded = torch.jit.load(results["torchscript"]["path"])
        assert loaded is not None

    def test_export_onnx_with_validation(self, resnet18_model, resnet18_input, tmp_path):
        """ONNX export with validate=True produces a validation dict with shape_match and values_match."""
        pipeline = ExportPipeline(
            resnet18_model, resnet18_input, task="classification"
        )
        results = pipeline.export(["onnx"], str(tmp_path), validate=True)

        assert results["onnx"]["status"] == "success"
        validation = results["onnx"]["validation"]
        assert validation is not None
        assert validation.get("shape_match") is True
        assert validation.get("values_match") is True

    def test_export_torchscript_with_validation(
        self, resnet18_model, resnet18_input, tmp_path
    ):
        """TorchScript export with validate=True produces a validation dict with shape_match=True.

        Note: TorchScript tracing for ResNet18 can introduce small numerical
        differences (max_diff ~ 0.004) that exceed the validator's default
        atol=1e-5, so ``values_match`` may be False.  We verify that
        shape_match is True and max_diff is within a reasonable tolerance.
        """
        pipeline = ExportPipeline(
            resnet18_model, resnet18_input, task="classification"
        )
        results = pipeline.export(["torchscript"], str(tmp_path), validate=True)

        assert results["torchscript"]["status"] == "success"
        validation = results["torchscript"]["validation"]
        assert validation is not None
        assert validation.get("shape_match") is True
        # TorchScript tracing may introduce small numerical diffs; verify
        # max_diff is within a reasonable tolerance (< 0.01).
        assert validation.get("max_diff") is not None
        assert validation["max_diff"] < 0.01


# ------------------------------------------------------------------ #
# 2. YOLOv8n exports via UltralyticsAdapter
# ------------------------------------------------------------------ #


@pytest.mark.real_model
class TestYOLOv8nUltralyticsExport:
    """Export YOLOv8n through the UltralyticsAdapter."""

    def test_pipeline_detects_yolo_as_pytorch(self, yolov8n_model):
        """ExportPipeline classifies a real YOLO model as 'pytorch'.

        Ultralytics YOLO models inherit from ``torch.nn.Module``, so the
        pipeline's type-detection logic classifies them as ``"pytorch"``
        (``_is_pytorch_model`` matches first).  The ``UltralyticsAdapter``
        is intended to be used directly for YOLO-specific export workflows.
        """
        pipeline = ExportPipeline(yolov8n_model, None)
        assert pipeline.model_type == "pytorch"

    def test_export_onnx_via_adapter(self, yolov8n_model):
        """UltralyticsAdapter.export produces an ONNX file that exists on disk."""
        adapter = UltralyticsAdapter()
        path = adapter.export(yolov8n_model, "onnx")
        assert path is not None
        assert os.path.exists(path)

    def test_export_torchscript_via_adapter(self, yolov8n_model):
        """UltralyticsAdapter.export produces a TorchScript file that exists on disk."""
        adapter = UltralyticsAdapter()
        path = adapter.export(yolov8n_model, "torchscript")
        assert path is not None
        assert os.path.exists(path)

    def test_export_all_basic_formats(self, yolov8n_model, tmp_path):
        """export_all with onnx and torchscript returns success for both formats."""
        adapter = UltralyticsAdapter()
        results = adapter.export_all(
            yolov8n_model, ["onnx", "torchscript"], str(tmp_path)
        )

        assert "onnx" in results
        assert "torchscript" in results
        assert results["onnx"]["status"] == "success"
        assert results["torchscript"]["status"] == "success"
        assert os.path.exists(results["onnx"]["path"])
        assert os.path.exists(results["torchscript"]["path"])

    def test_adapter_unsupported_format_raises(self, yolov8n_model):
        """Requesting an unsupported format from the adapter raises ValueError."""
        adapter = UltralyticsAdapter()
        with pytest.raises(ValueError, match="Unsupported Ultralytics format"):
            adapter.export(yolov8n_model, "fake_format")


# ------------------------------------------------------------------ #
# 3. ResNet18 pipeline baseline generation
# ------------------------------------------------------------------ #


@pytest.mark.real_model
class TestResNet18PipelineBaseline:
    """Verify baseline output generation for ResNet18."""

    def test_baseline_generated(self, resnet18_model, resnet18_input):
        """ExportPipeline generates a non-None baseline_output for a PyTorch model."""
        pipeline = ExportPipeline(resnet18_model, resnet18_input)
        assert pipeline.baseline_output is not None

    def test_baseline_shape(self, resnet18_model, resnet18_input):
        """The baseline output has shape (1, 1000) matching ResNet18 ImageNet classes."""
        pipeline = ExportPipeline(resnet18_model, resnet18_input)
        assert pipeline.baseline_output.shape == (1, 1000)

    def test_model_type_is_pytorch(self, resnet18_model, resnet18_input):
        """ResNet18 (nn.Module) is classified as 'pytorch'."""
        pipeline = ExportPipeline(resnet18_model, resnet18_input)
        assert pipeline.model_type == "pytorch"
