# Copyright (C) 2025 Embedl AB

"""
Test cases for the embedl-hub CLI quantize command.

This test suite covers the refactored quantize CLI with the following structure:
- Main quantize command that defaults to TFLite quantization
- qai-hub subcommand for ONNX quantization using Qualcomm AI Hub
- embedl subcommand for TFLite quantization

Test coverage includes:
- Command invocation with various parameter combinations
- Error handling for missing arguments and invalid inputs
- Proper integration between main command and subcommands
- Mocking of external dependencies (quantizers, file operations, etc.)
"""

from pathlib import Path
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
from typer.testing import CliRunner

from embedl_hub.cli.quantize import quantize_cli


class MockQuantizationResult:
    """Mock result object returned by the quantizer."""

    def __init__(self, model_path: Path):
        self.model_path = model_path


@pytest.fixture
def runner():
    """Provide a CLI runner for testing."""
    return CliRunner()


@pytest.fixture
def default_ctx_obj():
    """Provide a context object with config containing a project and experiment."""
    return {
        "config": {
            "project_name": "test_project",
            "experiment_name": "test_experiment",
        },
        "state": {},
    }


@pytest.fixture
def ctx_obj(default_ctx_obj):
    """Provide a new CLI context object."""
    return default_ctx_obj.copy()


@pytest.fixture
def temp_model_file(tmp_path):
    """Create a temporary model file for testing."""
    model_file = tmp_path / "test_model.onnx"
    model_file.write_text("dummy onnx model content")
    return model_file


@pytest.fixture
def temp_tflite_model_file(tmp_path):
    """Create a temporary TFLite model file for testing."""
    model_file = tmp_path / "test_model.tflite"
    model_file.write_text("dummy tflite model content")
    return model_file


@pytest.fixture
def temp_data_dir(tmp_path):
    """Create a temporary data directory for testing."""
    data_dir = tmp_path / "cal_data"
    data_dir.mkdir()
    # Create some dummy .npy files
    for i in range(3):
        dummy_data = np.random.rand(1, 3, 224, 224)
        np.save(data_dir / f"sample_{i}.npy", dummy_data)
    return data_dir


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.utils.onnx_utils.maybe_package_onnx_folder_to_file")
@patch("embedl_hub.core.quantization.quantize.QAIHubQuantizer")
def test_qai_hub_quantize_command_minimal_args(
    mock_quantizer_class,
    mock_package_onnx,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_model_file,
):
    """Test qai-hub quantize command with minimal required arguments."""
    # Arrange
    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer

    expected_output = temp_model_file.with_suffix(".quantized.onnx")
    mock_result = MockQuantizationResult(expected_output)
    mock_quantizer.quantize.return_value = mock_result

    mock_package_onnx.return_value = temp_model_file

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(temp_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0
    mock_assert_api_config.assert_called_once()
    mock_require_ctx.assert_called_once_with(ctx_obj["config"])

    # Check that quantizer was created with default values
    mock_quantizer_class.assert_called_once_with(
        data_path=None, num_samples=500
    )

    # Check that quantize was called correctly
    mock_quantizer.quantize.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_model_file,
        output_path=expected_output,
    )

    # Check console output
    mock_console.print.assert_any_call(
        f"[yellow]No output file specified, using default: {expected_output}[/]"
    )
    mock_console.print.assert_any_call(
        "[yellow]No data path specified, generating random calibration data.[/]"
    )
    mock_console.print.assert_any_call(
        f"[green]✓ Quantized model saved to {expected_output}[/]"
    )


@pytest.mark.parametrize(
    "use_short_flags",
    [False, True],
    ids=["long_flags", "short_flags"],
)
@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.utils.onnx_utils.maybe_package_onnx_folder_to_file")
@patch("embedl_hub.core.quantization.quantize.QAIHubQuantizer")
def test_qai_hub_quantize_command_all_args(
    mock_quantizer_class,
    mock_package_onnx,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_model_file,
    temp_data_dir,
    tmp_path,
    use_short_flags,
):
    """Test qai-hub quantize command with all arguments provided using both long and short flags."""
    # Arrange
    output_file = tmp_path / "custom_output.quantized.onnx"
    num_samples = 100

    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer

    mock_result = MockQuantizationResult(output_file)
    mock_quantizer.quantize.return_value = mock_result

    mock_package_onnx.return_value = temp_model_file

    # Choose flags based on parameter
    if use_short_flags:
        args = [
            "qai-hub",
            "-m",
            str(temp_model_file),
            "-d",
            str(temp_data_dir),
            "-o",
            str(output_file),
            "-n",
            str(num_samples),
        ]
    else:
        args = [
            "qai-hub",
            "--model",
            str(temp_model_file),
            "--data",
            str(temp_data_dir),
            "--output-file",
            str(output_file),
            "--num-samples",
            str(num_samples),
        ]

    # Act
    result = runner.invoke(quantize_cli, args, obj=ctx_obj)

    # Assert
    assert result.exit_code == 0
    mock_assert_api_config.assert_called_once()
    mock_require_ctx.assert_called_once_with(ctx_obj["config"])

    # Check that quantizer was created with provided values
    mock_quantizer_class.assert_called_once_with(
        data_path=temp_data_dir, num_samples=num_samples
    )

    # Check that quantize was called correctly
    mock_quantizer.quantize.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_model_file,
        output_path=output_file,
    )

    # Should not show default messages when all args provided
    for call in mock_console.print.call_args_list:
        assert "No output file specified" not in str(call)
        assert "No data path specified" not in str(call)

    # Should show success message
    mock_console.print.assert_any_call(
        f"[green]✓ Quantized model saved to {output_file}[/]"
    )


@patch("embedl_hub.cli.utils.assert_api_config")
def test_qai_hub_quantize_command_missing_model_arg(
    mock_assert_api_config, runner, ctx_obj
):
    """Test that qai-hub quantize command fails when model argument is missing."""
    # Act
    result = runner.invoke(quantize_cli, ["qai-hub"], obj=ctx_obj)

    # Assert
    assert result.exit_code != 0
    mock_assert_api_config.assert_called_once()
    # With no_args_is_help=True, the command shows help instead of an error message
    assert "Usage: quantize qai-hub [OPTIONS]" in result.output


@patch("embedl_hub.cli.utils.assert_api_config")
def test_qai_hub_quantize_command_api_config_assertion_error(
    mock_assert_api_config, runner, ctx_obj, temp_model_file
):
    """Test that qai-hub quantize command fails when API config assertion fails."""
    # Arrange
    mock_assert_api_config.side_effect = Exception("API config error")

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(temp_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code != 0
    mock_assert_api_config.assert_called_once()


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
def test_qai_hub_quantize_command_context_requirement_error(
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_model_file,
):
    """Test that qai-hub quantize command fails when context requirement fails."""
    # Arrange
    mock_require_ctx.side_effect = Exception("Context not initialized")

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(temp_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code != 0
    mock_assert_api_config.assert_called_once()
    mock_require_ctx.assert_called_once_with(ctx_obj["config"])


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.utils.onnx_utils.maybe_package_onnx_folder_to_file")
@patch("embedl_hub.core.quantization.quantize.QAIHubQuantizer")
def test_qai_hub_quantize_command_quantizer_error(
    mock_quantizer_class,
    mock_package_onnx,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_model_file,
):
    """Test that qai-hub quantize command handles quantizer errors gracefully."""
    # Arrange
    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer
    mock_quantizer.quantize.side_effect = Exception("Quantization failed")

    mock_package_onnx.return_value = temp_model_file

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(temp_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code != 0
    mock_quantizer.quantize.assert_called_once()


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.utils.onnx_utils.maybe_package_onnx_folder_to_file")
@patch("embedl_hub.core.quantization.quantize.QAIHubQuantizer")
def test_qai_hub_quantize_command_with_model_directory(
    mock_quantizer_class,
    mock_package_onnx,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    tmp_path,
):
    """Test qai-hub quantize command with a model directory instead of single file."""
    # Arrange
    model_dir = tmp_path / "model_dir"
    model_dir.mkdir()
    (model_dir / "model.onnx").write_text("dummy model")
    (model_dir / "model.data").write_text("dummy data")

    packaged_model = tmp_path / "packaged_model.onnx"
    mock_package_onnx.return_value = packaged_model

    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer

    expected_output = model_dir.with_suffix(".quantized")
    mock_result = MockQuantizationResult(expected_output)
    mock_quantizer.quantize.return_value = mock_result

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(model_dir)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0
    # Verify that maybe_package_onnx_folder_to_file was called with model_dir
    mock_package_onnx.assert_called_once()
    args, _ = mock_package_onnx.call_args
    assert Path(args[0]) == model_dir


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.utils.onnx_utils.maybe_package_onnx_folder_to_file")
@patch("embedl_hub.core.quantization.quantize.QAIHubQuantizer")
@patch("tempfile.TemporaryDirectory")
def test_qai_hub_quantize_command_temporary_directory_usage(
    mock_temp_dir,
    mock_quantizer_class,
    mock_package_onnx,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_model_file,
):
    """Test that TemporaryDirectory is used correctly in qai-hub quantize command."""
    # Arrange
    mock_temp_dir_instance = MagicMock()
    mock_temp_dir_instance.__enter__.return_value = "/tmp/test_dir"
    mock_temp_dir.return_value = mock_temp_dir_instance

    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer

    expected_output = temp_model_file.with_suffix(".quantized")
    mock_result = MockQuantizationResult(expected_output)
    mock_quantizer.quantize.return_value = mock_result

    mock_package_onnx.return_value = temp_model_file

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(temp_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0
    mock_temp_dir.assert_called_once()
    mock_temp_dir_instance.__enter__.assert_called_once()
    mock_temp_dir_instance.__exit__.assert_called_once()

    # Verify maybe_package_onnx_folder_to_file was called with tmpdir
    mock_package_onnx.assert_called_once_with(temp_model_file, "/tmp/test_dir")


def test_qai_hub_quantize_command_invalid_model_path(runner, ctx_obj):
    """Test qai-hub quantize command with non-existent model path."""
    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", "/non/existent/path.onnx"],
        obj=ctx_obj,
    )

    # Assert - The command should still try to run but may fail during execution
    # The exact behavior depends on how the underlying functions handle invalid paths
    # This test primarily ensures the CLI parsing works correctly
    assert "--model" in result.output or result.exit_code != 0


def test_qai_hub_quantize_command_invalid_num_samples(
    runner, ctx_obj, temp_model_file
):
    """Test qai-hub quantize command with invalid num_samples value."""
    # Act - Try with negative number
    result = runner.invoke(
        quantize_cli,
        [
            "qai-hub",
            "--model",
            str(temp_model_file),
            "--num-samples",
            "-1",
        ],
        obj=ctx_obj,
    )

    # Assert - Should either reject negative values or handle them gracefully
    # The exact behavior depends on typer's validation
    assert result.exit_code is not None


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.utils.onnx_utils.maybe_package_onnx_folder_to_file")
@patch("embedl_hub.core.quantization.quantize.QAIHubQuantizer")
def test_qai_hub_quantize_command_default_output_file_generation(
    mock_quantizer_class,
    mock_package_onnx,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    tmp_path,
):
    """Test that default output file is generated correctly based on model name for qai-hub quantize command."""
    # Arrange
    model_file = tmp_path / "my_model.onnx"
    model_file.write_text("dummy model")
    expected_output = tmp_path / "my_model.quantized.onnx"

    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer

    mock_result = MockQuantizationResult(expected_output)
    mock_quantizer.quantize.return_value = mock_result

    mock_package_onnx.return_value = model_file

    # Act
    result = runner.invoke(
        quantize_cli,
        ["qai-hub", "--model", str(model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0

    # Verify the default output file was used
    mock_quantizer.quantize.assert_called_once()
    _, kwargs = mock_quantizer.quantize.call_args
    assert kwargs["output_path"] == expected_output

    # Verify console message about default output
    mock_console.print.assert_any_call(
        f"[yellow]No output file specified, using default: {expected_output}[/]"
    )


# Tests for the main quantize command callback


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
def test_quantize_main_command_with_model_invokes_tflite_quantize(
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_model_file,
):
    """Test that main quantize command with --model invokes the tflite quantize command."""
    with patch(
        "embedl_hub.cli.quantize.tflite_quantize_command"
    ) as mock_tflite_command:
        # Act
        result = runner.invoke(
            quantize_cli,
            ["--model", str(temp_model_file)],
            obj=ctx_obj,
        )

        # Assert
        assert result.exit_code == 0
        mock_assert_api_config.assert_called_once()
        mock_require_ctx.assert_called_once_with(ctx_obj["config"])
        mock_tflite_command.assert_called_once()


@patch("embedl_hub.cli.utils.assert_api_config")
def test_quantize_main_command_missing_model_fails(
    mock_assert_api_config, runner, ctx_obj
):
    """Test that main quantize command without --model fails."""
    # Act
    result = runner.invoke(quantize_cli, [], obj=ctx_obj)

    # Assert
    assert result.exit_code != 0
    mock_assert_api_config.assert_called_once()
    assert (
        "Please specify a model to compile using --model" in result.output
        or "Error" in result.output
    )


# Tests for the TFLite embedl subcommand


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.quantization.TFLiteQuantizer")
def test_tflite_quantize_command_minimal_args(
    mock_quantizer_class,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_tflite_model_file,
):
    """Test TFLite quantize command with minimal required arguments."""
    # Arrange
    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer
    mock_quantizer_class.validate_core_args = MagicMock()

    expected_output = temp_tflite_model_file.with_suffix(".quantized.tflite")
    mock_result = MockQuantizationResult(expected_output)
    mock_quantizer.quantize.return_value = mock_result

    # Act
    result = runner.invoke(
        quantize_cli,
        ["embedl", "--model", str(temp_tflite_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0
    mock_assert_api_config.assert_called_once()
    mock_require_ctx.assert_called_once_with(ctx_obj["config"])

    # Check that validator was called
    mock_quantizer_class.validate_core_args.assert_called_once_with(
        temp_tflite_model_file, expected_output
    )

    # Check that quantizer was created
    mock_quantizer_class.assert_called_once()

    # Check that quantize was called correctly
    mock_quantizer.quantize.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_tflite_model_file,
        output_path=expected_output,
    )

    # Check console output
    mock_console.print.assert_any_call(
        f"[yellow]No output file specified, using default: {expected_output}[/]"
    )
    mock_console.print.assert_any_call(
        f"[green]✓ Quantized model saved to {expected_output}[/]"
    )


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.quantization.TFLiteQuantizer")
def test_tflite_quantize_command_with_output_file(
    mock_quantizer_class,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_tflite_model_file,
    tmp_path,
):
    """Test TFLite quantize command with custom output file."""
    # Arrange
    custom_output = tmp_path / "custom_quantized_model.tflite"
    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer
    mock_quantizer_class.validate_core_args = MagicMock()

    mock_result = MockQuantizationResult(custom_output)
    mock_quantizer.quantize.return_value = mock_result

    # Act
    result = runner.invoke(
        quantize_cli,
        [
            "embedl",
            "--model",
            str(temp_tflite_model_file),
            "--output-file",
            str(custom_output),
        ],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0

    # Check that validator was called with custom output
    mock_quantizer_class.validate_core_args.assert_called_once_with(
        temp_tflite_model_file, custom_output
    )

    # Check that quantize was called with custom output
    mock_quantizer.quantize.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_tflite_model_file,
        output_path=custom_output,
    )

    # Should not show default output message when custom output provided
    for call in mock_console.print.call_args_list:
        assert "No output file specified" not in str(call)

    # Should show success message
    mock_console.print.assert_any_call(
        f"[green]✓ Quantized model saved to {custom_output}[/]"
    )


@patch("embedl_hub.cli.utils.assert_api_config")
def test_tflite_quantize_command_missing_model_arg(
    mock_assert_api_config, runner, ctx_obj
):
    """Test that TFLite quantize command fails when model argument is missing."""
    # Act
    result = runner.invoke(quantize_cli, ["embedl"], obj=ctx_obj)

    # Assert
    assert result.exit_code != 0
    mock_assert_api_config.assert_called_once()
    # With no_args_is_help=True, the command shows help instead of an error message
    assert "Usage: quantize embedl [OPTIONS]" in result.output


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.hub_logging.console")
@patch("embedl_hub.core.quantization.TFLiteQuantizer")
def test_tflite_quantize_command_quantizer_error(
    mock_quantizer_class,
    mock_console,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_tflite_model_file,
):
    """Test that TFLite quantize command handles quantizer errors gracefully."""
    # Arrange
    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer
    mock_quantizer_class.validate_core_args = MagicMock()
    mock_quantizer.quantize.side_effect = Exception(
        "TFLite quantization failed"
    )

    # Act
    result = runner.invoke(
        quantize_cli,
        ["embedl", "--model", str(temp_tflite_model_file)],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code != 0
    mock_quantizer.quantize.assert_called_once()


@patch("embedl_hub.cli.utils.assert_api_config")
@patch("embedl_hub.core.context.require_initialized_ctx")
@patch("embedl_hub.core.quantization.TFLiteQuantizer")
def test_tflite_quantize_command_with_data_path(
    mock_quantizer_class,
    mock_require_ctx,
    mock_assert_api_config,
    runner,
    ctx_obj,
    temp_tflite_model_file,
    temp_data_dir,
):
    """Test TFLite quantize command with custom data path for calibration."""
    # Arrange
    mock_quantizer = MagicMock()
    mock_quantizer_class.return_value = mock_quantizer
    mock_quantizer_class.validate_core_args = MagicMock()

    expected_output = temp_tflite_model_file.with_suffix(".quantized.tflite")
    mock_result = MockQuantizationResult(expected_output)
    mock_quantizer.quantize.return_value = mock_result

    # Act
    result = runner.invoke(
        quantize_cli,
        [
            "embedl",
            "--model",
            str(temp_tflite_model_file),
            "--data",
            str(temp_data_dir),
        ],
        obj=ctx_obj,
    )

    # Assert
    assert result.exit_code == 0
    mock_assert_api_config.assert_called_once()
    mock_require_ctx.assert_called_once_with(ctx_obj["config"])

    # Check that quantizer was created with data_path parameter
    mock_quantizer_class.assert_called_once_with(
        num_samples=500,  # default value
        data_path=temp_data_dir,
    )

    # Check that quantize was called correctly
    mock_quantizer.quantize.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_tflite_model_file,
        output_path=expected_output,
    )


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