# Copyright (C) 2025 Embedl AB

"""
Tests for compilation of ONNX models to TensorFlow/TFLite format.
"""

from pathlib import Path
from unittest.mock import Mock, mock_open, patch

import pytest

from embedl_hub.core.compile.onnx_to_tf import compile_onnx_to_tflite


@patch('builtins.open', new_callable=mock_open)
@patch('embedl_hub.core.compile.onnx_to_tf.tf.lite.TFLiteConverter')
@patch('embedl_hub.core.compile.onnx_to_tf.onnx2tf')
@patch('embedl_hub.core.compile.onnx_to_tf.TemporaryDirectory')
def test_compile_onnx_to_tflite_with_output_model_path(
    mock_tmpdir, mock_onnx2tf, mock_converter_class, mock_file_open
):
    """Test conversion with specified output folder."""
    # Setup mocks
    mock_tmpdir.return_value.__enter__.return_value = "/tmp/test"
    mock_converter = Mock()
    mock_converter.convert.return_value = b"fake tflite data"
    mock_converter_class.from_saved_model.return_value = mock_converter

    onnx_path = Path("/path/to/model.onnx")
    output_model_path = Path("/path/to/model.tflite")

    result = compile_onnx_to_tflite(onnx_path, output_model_path)

    # Verify onnx2tf.convert was called correctly
    mock_onnx2tf.convert.assert_called_once_with(
        input_onnx_file_path=onnx_path,
        output_folder_path="/tmp/test",
        output_signaturedefs=True,
    )

    # Verify TFLiteConverter was used correctly
    mock_converter_class.from_saved_model.assert_called_once_with("/tmp/test")
    mock_converter.convert.assert_called_once()

    # Verify file was written
    mock_file_open.assert_called_once_with(output_model_path, "wb")
    mock_file_open().write.assert_called_once_with(b"fake tflite data")

    # Verify return value
    assert result.model_path == output_model_path


@patch('builtins.open', new_callable=mock_open)
@patch('embedl_hub.core.compile.onnx_to_tf.tf.lite.TFLiteConverter')
@patch('embedl_hub.core.compile.onnx_to_tf.onnx2tf')
@patch('embedl_hub.core.compile.onnx_to_tf.TemporaryDirectory')
def test_compile_onnx_to_tflite_without_output_model_path(
    mock_tmpdir, mock_onnx2tf, mock_converter_class, mock_file_open
):
    """Test conversion without specified output folder."""
    # Setup mocks
    mock_tmpdir.return_value.__enter__.return_value = "/tmp/test"
    mock_converter = Mock()
    mock_converter.convert.return_value = b"fake tflite data"
    mock_converter_class.from_saved_model.return_value = mock_converter

    onnx_path = Path("/path/to/model.onnx")
    expected_output_path = onnx_path.with_suffix(".tflite")

    result = compile_onnx_to_tflite(onnx_path)

    # Verify onnx2tf.convert was called correctly
    mock_onnx2tf.convert.assert_called_once_with(
        input_onnx_file_path=onnx_path,
        output_folder_path="/tmp/test",
        output_signaturedefs=True,
    )

    # Verify TFLiteConverter was used correctly
    mock_converter_class.from_saved_model.assert_called_once_with("/tmp/test")
    mock_converter.convert.assert_called_once()

    # Verify file was written
    mock_file_open.assert_called_once_with(expected_output_path, "wb")
    mock_file_open().write.assert_called_once_with(b"fake tflite data")

    # Verify return value
    assert result.model_path == expected_output_path


@patch('builtins.open', new_callable=mock_open)
@patch('embedl_hub.core.compile.onnx_to_tf.tf.lite.TFLiteConverter')
@patch('embedl_hub.core.compile.onnx_to_tf.onnx2tf')
@patch('embedl_hub.core.compile.onnx_to_tf.TemporaryDirectory')
def test_compile_onnx_to_tflite_with_none_output_model_path(
    mock_tmpdir, mock_onnx2tf, mock_converter_class, mock_file_open
):
    """Test conversion with explicitly None output folder."""
    # Setup mocks
    mock_tmpdir.return_value.__enter__.return_value = "/tmp/test"
    mock_converter = Mock()
    mock_converter.convert.return_value = b"fake tflite data"
    mock_converter_class.from_saved_model.return_value = mock_converter

    onnx_path = Path("/path/to/model.onnx")
    expected_output_path = onnx_path.with_suffix(".tflite")

    result = compile_onnx_to_tflite(onnx_path, None)

    # Verify onnx2tf.convert was called correctly
    mock_onnx2tf.convert.assert_called_once_with(
        input_onnx_file_path=onnx_path,
        output_folder_path="/tmp/test",
        output_signaturedefs=True,
    )

    # Verify TFLiteConverter was used correctly
    mock_converter_class.from_saved_model.assert_called_once_with("/tmp/test")
    mock_converter.convert.assert_called_once()

    # Verify file was written
    mock_file_open.assert_called_once_with(expected_output_path, "wb")
    mock_file_open().write.assert_called_once_with(b"fake tflite data")

    # Verify return value
    assert result.model_path == expected_output_path


if __name__ == "__main__":
    pytest.main([__file__])
