"""Pipeline edge case tests — empty formats, numpy input, custom options, unusual models."""

from __future__ import annotations

import os

import numpy as np
import onnx
import torch

from matrice_export import ExportPipeline

# ------------------------------------------------------------------ #
# 1. Empty format list
# ------------------------------------------------------------------ #


class TestExportEmptyFormatList:
    def test_empty_formats_returns_empty_dict(self, dummy_model, sample_input, tmp_path):
        """export([]) returns an empty dict without error."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export([], str(tmp_path), validate=False)
        assert results == {}


# ------------------------------------------------------------------ #
# 2. Duplicate formats
# ------------------------------------------------------------------ #


class TestExportDuplicateFormats:
    def test_duplicate_format_returns_last_result(self, dummy_model, sample_input, tmp_path):
        """Requesting the same format twice doesn't crash; result dict has one entry."""
        pipeline = ExportPipeline(dummy_model, sample_input)
        results = pipeline.export(["onnx", "onnx"], str(tmp_path), validate=False)
        # Dict keys are unique, so second export overwrites first
        assert "onnx" in results
        assert results["onnx"]["status"] == "success"


# ------------------------------------------------------------------ #
# 3. Numpy input
# ------------------------------------------------------------------ #


class TestExportWithNumpyInput:
    def test_numpy_sample_input_works(self, dummy_model, tmp_path):
        """Passing sample_input as numpy array (not Tensor) works."""
        np_input = np.random.randn(1, 10).astype(np.float32)
        pipeline = ExportPipeline(dummy_model, np_input)
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        assert results["onnx"]["status"] == "success"

    def test_numpy_baseline_generated(self, dummy_model):
        """Baseline is generated even with numpy sample_input."""
        np_input = np.random.randn(1, 10).astype(np.float32)
        pipeline = ExportPipeline(dummy_model, np_input)
        assert pipeline.baseline_output is not None
        assert pipeline.baseline_output.shape == (1, 5)


# ------------------------------------------------------------------ #
# 4. Conv model multi-format with validation
# ------------------------------------------------------------------ #


class TestConvModelMultiFormat:
    def test_conv_onnx_and_torchscript_with_validation(self, conv_model, image_input, tmp_path):
        """Conv model exports both onnx and torchscript successfully."""
        pipeline = ExportPipeline(conv_model, image_input, task="classification")
        results = pipeline.export(["onnx", "torchscript"], str(tmp_path), validate=False)
        assert results["onnx"]["status"] == "success"
        assert results["torchscript"]["status"] == "success"
        assert os.path.isfile(results["onnx"]["path"])
        assert os.path.isfile(results["torchscript"]["path"])


# ------------------------------------------------------------------ #
# 5. Custom opset
# ------------------------------------------------------------------ #


class TestExportWithCustomOpset:
    def test_opset_13(self, dummy_model, sample_input, tmp_path):
        """Export with opset=13 produces a valid ONNX file with opset 13."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False, opset=13)
        assert results["onnx"]["status"] == "success"

        model = onnx.load(results["onnx"]["path"])
        onnx.checker.check_model(model)

    def test_opset_17(self, dummy_model, sample_input, tmp_path):
        """Export with opset=17 produces a valid ONNX file."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False, opset=17)
        assert results["onnx"]["status"] == "success"


# ------------------------------------------------------------------ #
# 6. Baseline with HuggingFace-style output
# ------------------------------------------------------------------ #


class _FakeHFOutput:
    """Mimics a HuggingFace model output with .logits attribute."""
    def __init__(self, logits):
        self.logits = logits


class _HFStyleModel(torch.nn.Module):
    """Model that returns an object with .logits."""
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)

    def forward(self, x):
        return _FakeHFOutput(self.linear(x))


class TestPipelineHuggingFaceOutput:
    def test_hf_style_baseline_extraction(self):
        """Baseline extraction handles model outputs with .logits attribute."""
        model = _HFStyleModel()
        sample = torch.randn(1, 10)
        pipeline = ExportPipeline(model, sample)
        assert pipeline.baseline_output is not None
        assert pipeline.baseline_output.shape == (1, 5)


# ------------------------------------------------------------------ #
# 7. Tuple output model
# ------------------------------------------------------------------ #


class _TupleOutputModel(torch.nn.Module):
    """Model that returns a tuple of tensors."""
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)

    def forward(self, x):
        out = self.linear(x)
        return out, out * 2


class TestPipelineTupleOutput:
    def test_tuple_output_baseline_uses_first(self):
        """Baseline extraction takes first element from tuple output."""
        model = _TupleOutputModel()
        sample = torch.randn(1, 10)
        pipeline = ExportPipeline(model, sample)
        assert pipeline.baseline_output is not None
        assert pipeline.baseline_output.shape == (1, 5)


# ------------------------------------------------------------------ #
# 8. Registry auto-registration
# ------------------------------------------------------------------ #


class TestRegistryAutoRegister:
    def test_exporters_populated(self):
        """After auto_register, EXPORTERS contains 'onnx' and 'torchscript'."""
        # Creating a pipeline triggers _auto_register
        assert "onnx" in ExportPipeline.EXPORTERS
        assert "torchscript" in ExportPipeline.EXPORTERS

    def test_validators_populated(self):
        """After auto_register, VALIDATORS contains 'onnx' and 'torchscript'."""
        assert "onnx" in ExportPipeline.VALIDATORS
        assert "torchscript" in ExportPipeline.VALIDATORS

    def test_all_11_formats_registered(self):
        """All 11 format exporters are registered."""
        expected = {
            "onnx", "torchscript", "openvino", "tensorrt",
            "coreml", "saved_model", "pb", "tflite",
            "edgetpu", "tfjs", "paddle",
        }
        actual = set(ExportPipeline.EXPORTERS.keys())
        assert expected.issubset(actual), f"Missing: {expected - actual}"


# ------------------------------------------------------------------ #
# 9. Available formats metadata
# ------------------------------------------------------------------ #


class TestAvailableFormatsMetadata:
    def test_tensorrt_requires_gpu(self):
        """TensorRT format is flagged as requires_gpu=True."""
        formats = ExportPipeline.available_formats()
        trt = [f for f in formats if f["name"] == "tensorrt"]
        assert len(trt) == 1
        assert trt[0]["requires_gpu"] is True

    def test_onnx_has_validator(self):
        """ONNX format is flagged as has_validator=True."""
        formats = ExportPipeline.available_formats()
        onnx_fmt = [f for f in formats if f["name"] == "onnx"]
        assert len(onnx_fmt) == 1
        assert onnx_fmt[0]["has_validator"] is True

    def test_torchscript_has_validator(self):
        """TorchScript format is flagged as has_validator=True."""
        formats = ExportPipeline.available_formats()
        ts = [f for f in formats if f["name"] == "torchscript"]
        assert len(ts) == 1
        assert ts[0]["has_validator"] is True

    def test_onnx_does_not_require_gpu(self):
        """ONNX export does not require GPU."""
        formats = ExportPipeline.available_formats()
        onnx_fmt = [f for f in formats if f["name"] == "onnx"]
        assert onnx_fmt[0]["requires_gpu"] is False
