"""Tests for the core ExportPipeline — model detection, export dispatch, and tracker integration."""

from __future__ import annotations

import os
from unittest.mock import MagicMock

import numpy as np
import pytest

from matrice_export import ExportPipeline, ExportTracker

# Import fixtures / helpers from conftest
from .conftest import MockActionTracker, MockYOLO

# ------------------------------------------------------------------ #
# 1-3. Model type detection
# ------------------------------------------------------------------ #

class TestModelTypeDetection:
    """Verify that the pipeline correctly classifies the three supported model types."""

    def test_pytorch_model_detected(self, dummy_model, sample_input):
        """A torch.nn.Module is classified as 'pytorch'."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        assert pipeline.model_type == "pytorch"

    def test_onnx_path_detected(self, onnx_model_path):
        """A string ending in .onnx that exists on disk is classified as 'onnx'."""
        pipeline = ExportPipeline(onnx_model_path)
        assert pipeline.model_type == "onnx"

    def test_ultralytics_model_detected(self, tmp_path):
        """An object with .export() and .predict() (but not nn.Module) is classified as 'ultralytics'."""
        mock_yolo = MockYOLO(export_path=str(tmp_path / "yolo.onnx"))
        pipeline = ExportPipeline(mock_yolo)
        assert pipeline.model_type == "ultralytics"


# ------------------------------------------------------------------ #
# 4. Unsupported model type
# ------------------------------------------------------------------ #

class TestUnsupportedModel:
    def test_unsupported_type_raises_type_error(self):
        """Passing an unsupported object (e.g. a plain int) raises TypeError."""
        with pytest.raises(TypeError, match="Unsupported model type"):
            ExportPipeline(12345)

    def test_dict_raises_type_error(self):
        """A dict is not a valid model type."""
        with pytest.raises(TypeError, match="Unsupported model type"):
            ExportPipeline({"weights": "path"})


# ------------------------------------------------------------------ #
# 5. Export ONNX from PyTorch
# ------------------------------------------------------------------ #

class TestExportOnnx:
    def test_export_onnx_success(self, dummy_model, sample_input, tmp_path):
        """Exporting a simple Linear model to ONNX produces a file with status 'success'."""
        pipeline = ExportPipeline(dummy_model, sample_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_onnx_file_extension(self, dummy_model, sample_input, tmp_path):
        """The exported ONNX file has a .onnx extension."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        assert results["onnx"]["path"].endswith(".onnx")


# ------------------------------------------------------------------ #
# 6. Export TorchScript from PyTorch
# ------------------------------------------------------------------ #

class TestExportTorchScript:
    def test_export_torchscript_success(self, dummy_model, sample_input, tmp_path):
        """Exporting a simple Linear model to TorchScript produces a file with status 'success'."""
        pipeline = ExportPipeline(dummy_model, sample_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"])

    def test_torchscript_file_extension(self, dummy_model, sample_input, tmp_path):
        """The exported TorchScript file has a .torchscript extension."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["torchscript"], str(tmp_path), validate=False)
        assert results["torchscript"]["path"].endswith(".torchscript")


# ------------------------------------------------------------------ #
# 7. Export multiple formats at once
# ------------------------------------------------------------------ #

class TestMultiFormatExport:
    def test_export_onnx_and_torchscript(self, dummy_model, sample_input, tmp_path):
        """Requesting both onnx and torchscript in a single call produces two results."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx", "torchscript"], str(tmp_path), validate=False)

        assert len(results) == 2
        assert results["onnx"]["status"] == "success"
        assert results["torchscript"]["status"] == "success"

    def test_multi_format_files_on_disk(self, dummy_model, sample_input, tmp_path):
        """Both exported files exist on disk after a multi-format export."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["onnx", "torchscript"], str(tmp_path), validate=False)

        for fmt in ["onnx", "torchscript"]:
            assert os.path.exists(results[fmt]["path"]), f"{fmt} file not found on disk"


# ------------------------------------------------------------------ #
# 8-9. Validate flag
# ------------------------------------------------------------------ #

class TestValidationFlag:
    def test_export_validate_true_no_validators(self, dummy_model, sample_input, tmp_path):
        """With validate=True but no validators registered, validation is None (no crash)."""
        # Clear validators to ensure none are registered
        saved = ExportPipeline.VALIDATORS.copy()
        ExportPipeline.VALIDATORS.clear()
        try:
            pipeline = ExportPipeline(dummy_model, sample_input)
            results = pipeline.export(["onnx"], str(tmp_path), validate=True)
            assert results["onnx"]["status"] == "success"
            assert results["onnx"]["validation"] is None
        finally:
            ExportPipeline.VALIDATORS.update(saved)

    def test_export_validate_false(self, dummy_model, sample_input, tmp_path):
        """With validate=False, validation is None regardless of registered validators."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        assert results["onnx"]["validation"] is None


# ------------------------------------------------------------------ #
# 10. Invalid format name
# ------------------------------------------------------------------ #

class TestInvalidFormat:
    def test_invalid_format_captured_as_error(self, dummy_model, sample_input, tmp_path):
        """An unregistered format name results in status='error' (ValueError captured)."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["nonexistent_format_xyz"], str(tmp_path), validate=False)

        assert results["nonexistent_format_xyz"]["status"] == "error"
        assert "No exporter registered" in results["nonexistent_format_xyz"]["error"]


# ------------------------------------------------------------------ #
# 11. Missing sample_input for PyTorch
# ------------------------------------------------------------------ #

class TestMissingSampleInput:
    def test_missing_sample_input_captured_as_error(self, dummy_model, tmp_path):
        """Exporting a PyTorch model without sample_input gives status='error'."""
        pipeline = ExportPipeline(dummy_model, sample_input=None)
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        assert results["onnx"]["status"] == "error"
        assert "sample_input is required" in results["onnx"]["error"]


# ------------------------------------------------------------------ #
# 12. available_formats()
# ------------------------------------------------------------------ #

class TestAvailableFormats:
    def test_available_formats_returns_list(self):
        """ExportPipeline.available_formats() returns a non-empty list of dicts."""
        formats = ExportPipeline.available_formats()
        assert isinstance(formats, list)
        assert len(formats) > 0

    def test_available_formats_contains_onnx(self):
        """The 'onnx' format is present in the available formats list."""
        names = [f["name"] for f in ExportPipeline.available_formats()]
        assert "onnx" in names

    def test_available_formats_dict_keys(self):
        """Each entry has the expected metadata keys."""
        for entry in ExportPipeline.available_formats():
            assert "name" in entry
            assert "suffix" in entry
            assert "requires_gpu" in entry
            assert "has_validator" in entry


# ------------------------------------------------------------------ #
# 13. ONNX path that doesn't exist
# ------------------------------------------------------------------ #

class TestOnnxPathNotFound:
    def test_nonexistent_onnx_path_raises_file_not_found(self):
        """Pointing to a non-existent .onnx file raises FileNotFoundError."""
        with pytest.raises(FileNotFoundError, match="ONNX model not found"):
            ExportPipeline("/tmp/does_not_exist_abc123.onnx")


# ------------------------------------------------------------------ #
# 14. Tracker integration
# ------------------------------------------------------------------ #

class TestTrackerIntegration:
    def test_tracker_methods_called(self, dummy_model, sample_input, tmp_path):
        """When a tracker is provided, its reporting methods are invoked."""
        tracker = MagicMock(spec=ExportTracker)
        pipeline = ExportPipeline(dummy_model, sample_input, tracker=tracker)
        pipeline.export(["onnx"], str(tmp_path), validate=False)

        tracker.report_pipeline_start.assert_called_once()
        tracker.report_pipeline_complete.assert_called_once()

    def test_tracker_export_success_called(self, dummy_model, sample_input, tmp_path):
        """On successful export, report_export_success is called."""
        tracker = MagicMock(spec=ExportTracker)
        pipeline = ExportPipeline(dummy_model, sample_input, tracker=tracker)
        pipeline.export(["onnx"], str(tmp_path), validate=False)

        tracker.report_export_success.assert_called_once()
        call_args = tracker.report_export_success.call_args
        assert call_args[0][0] == "onnx"  # first positional arg is format name

    def test_tracker_export_failure_called(self, dummy_model, tmp_path):
        """On failed export (missing sample_input), report_export_failure is called."""
        tracker = MagicMock(spec=ExportTracker)
        pipeline = ExportPipeline(dummy_model, sample_input=None, tracker=tracker)
        pipeline.export(["onnx"], str(tmp_path), validate=False)

        tracker.report_export_failure.assert_called_once()

    def test_export_tracker_with_action_tracker(self, dummy_model, sample_input, tmp_path):
        """ExportTracker backed by a MockActionTracker records status updates."""
        mock_at = MockActionTracker()
        tracker = ExportTracker(action_tracker=mock_at)
        pipeline = ExportPipeline(dummy_model, sample_input, tracker=tracker)
        pipeline.export(["onnx"], str(tmp_path), validate=False)

        # Expect at least pipeline ACK, pipeline START, format ACK, format RESULT, pipeline COMPLETE
        step_codes = [c[0] for c in mock_at.calls]
        assert "MDL_EXP_ACK" in step_codes
        assert "MDL_EXP_STR" in step_codes
        assert "MDL_EXP_CMPL" in step_codes


# ------------------------------------------------------------------ #
# 15. Export produces files on disk
# ------------------------------------------------------------------ #

class TestExportFilesOnDisk:
    def test_output_dir_created(self, dummy_model, sample_input, tmp_path):
        """Export creates the output directory if it does not exist."""
        out = str(tmp_path / "nested" / "exports")
        pipeline = ExportPipeline(dummy_model, sample_input)
        pipeline.export(["onnx"], out, validate=False)
        assert os.path.isdir(out)

    def test_exported_file_non_empty(self, dummy_model, sample_input, tmp_path):
        """The exported ONNX file is non-empty."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        file_size = os.path.getsize(results["onnx"]["path"])
        assert file_size > 0


# ------------------------------------------------------------------ #
# Additional: baseline output generation
# ------------------------------------------------------------------ #

class TestBaselineOutput:
    def test_baseline_generated_for_pytorch(self, dummy_model, sample_input):
        """Creating a pipeline with a PyTorch model and sample_input generates a baseline."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        assert pipeline.baseline_output is not None
        assert isinstance(pipeline.baseline_output, np.ndarray)

    def test_baseline_shape_matches_model_output(self, dummy_model, sample_input):
        """The baseline output shape matches the expected model output."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        # dummy_model: Linear(10,5) -> ReLU -> shape (1, 5)
        assert pipeline.baseline_output.shape == (1, 5)
