# Copyright (C) 2025 Embedl AB

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

from unittest.mock import Mock, patch

import pytest
import typer
from typer.testing import CliRunner

from embedl_hub.cli.compile import compile_cli

runner = CliRunner()


@pytest.fixture
def default_ctx_obj():
    """Provide an empty context object with config and state."""
    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("fake onnx model content")
    return model_file


@pytest.fixture
def temp_output_dir(tmp_path):
    """Create a temporary output directory for testing."""
    output_dir = tmp_path / "output"
    output_dir.mkdir()
    return output_dir


def test_compile_command_without_model_fails(ctx_obj, monkeypatch):
    """Test that compile command fails when no model is specified."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    result = runner.invoke(compile_cli, [], obj=ctx_obj)
    assert result.exit_code == 2
    # The error message appears in the output, not the exception string
    assert "Please specify a model to compile using --model" in result.output


@patch("embedl_hub.core.compile.ONNXToTFCompiler")
def test_compile_command_with_model_calls_embedl_default(
    mock_compiler_class,
    ctx_obj,
    temp_model_file,
    temp_output_dir,
    monkeypatch,
):
    """Test that compile command with model defaults to embedl subcommand."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_compiler = Mock()
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = [
        "--model",
        str(temp_model_file),
        "--output-file",
        str(temp_output_dir),
    ]
    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 0
    mock_compiler_class.validate_core_args.assert_called_once_with(
        model_path=temp_model_file, output_path=temp_output_dir
    )
    mock_compiler.compile.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_model_file,
        output_path=temp_output_dir,
    )


def test_compile_command_requires_initialized_context(ctx_obj, monkeypatch):
    """Test that compile command requires initialized context."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)

    def mock_require_initialized_ctx(config):
        raise typer.Exit(1)

    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx",
        mock_require_initialized_ctx,
    )

    result = runner.invoke(compile_cli, ["--model", "dummy.onnx"], obj=ctx_obj)
    assert result.exit_code == 1


def test_compile_command_requires_api_config(ctx_obj, monkeypatch):
    """Test that compile command requires API configuration."""

    def mock_assert_api_config():
        raise typer.Exit(1)

    monkeypatch.setattr(
        "embedl_hub.cli.utils.assert_api_config", mock_assert_api_config
    )
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    result = runner.invoke(compile_cli, ["--model", "dummy.onnx"], obj=ctx_obj)
    assert result.exit_code == 1


@patch("embedl_hub.core.compile.ONNXToTFCompiler")
def test_embedl_compile_success(
    mock_compiler_class,
    ctx_obj,
    temp_model_file,
    temp_output_dir,
    monkeypatch,
):
    """Test successful embedl compilation."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_compiler = Mock()
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = [
        "embedl",
        "--model",
        str(temp_model_file),
        "--output-file",
        str(temp_output_dir),
    ]
    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 0
    mock_compiler_class.validate_core_args.assert_called_once_with(
        model_path=temp_model_file, output_path=temp_output_dir
    )
    mock_compiler.compile.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_model_file,
        output_path=temp_output_dir,
    )


def test_embedl_compile_missing_model_fails(ctx_obj, monkeypatch):
    """Test that embedl compile fails when no model is specified."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )
    result = runner.invoke(compile_cli, ["embedl"], obj=ctx_obj)
    assert result.exit_code == 2
    assert "Missing option" in result.output or "Usage:" in result.output


@patch("embedl_hub.core.compile.ONNXToTFCompiler")
def test_embedl_compile_without_output_file(
    mock_compiler_class, ctx_obj, temp_model_file, monkeypatch
):
    """Test embedl compile without specifying output file."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_compiler = Mock()
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = ["embedl", "--model", str(temp_model_file)]
    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 0
    # Should call with None for output_path when not specified
    mock_compiler.compile.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_model_file,
        output_path=None,
    )


@patch("embedl_hub.core.compile.QAIHubCompiler")
@patch("embedl_hub.cli.utils.prepare_input_size")
@pytest.mark.parametrize(
    "runtime,quantize_io",
    [
        ("tflite", False),
        ("onnx", True),
        ("qnn", False),
    ],
)
def test_qai_hub_compile_success(
    mock_prepare_input_size,
    mock_compiler_class,
    ctx_obj,
    temp_model_file,
    temp_output_dir,
    monkeypatch,
    runtime,
    quantize_io,
):
    """Test successful QAI Hub compilation with different runtime and quantize_io combinations."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_prepare_input_size.return_value = [1, 3, 224, 224]
    mock_compiler = Mock()
    mock_result = Mock()
    mock_result.device = "Samsung Galaxy S24"
    mock_compiler.compile.return_value = mock_result
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = [
        "qai-hub",
        "--model",
        str(temp_model_file),
        "--output-file",
        str(temp_output_dir),
        "--device",
        "Samsung Galaxy S24",
        "--size",
        "1,3,224,224",
        "--runtime",
        runtime,
    ]
    if quantize_io:
        args.append("--quantize-io")

    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 0
    assert "✓ Compiled model for Samsung Galaxy S24" in result.output

    mock_prepare_input_size.assert_called_once_with("1,3,224,224")
    mock_compiler_class.assert_called_once_with(
        device="Samsung Galaxy S24",
        runtime=runtime,
        quantize_io=quantize_io,
        input_size=[1, 3, 224, 224],
    )
    mock_compiler_class.validate_core_args.assert_called_once_with(
        model_path=temp_model_file, output_path=temp_output_dir
    )
    mock_compiler.compile.assert_called_once_with(
        project_name="test_project",
        experiment_name="test_experiment",
        model_path=temp_model_file,
        output_path=temp_output_dir,
    )


@patch("embedl_hub.core.compile.QAIHubCompiler")
@patch("embedl_hub.cli.utils.prepare_input_size")
@patch("embedl_hub.cli.utils.console")
def test_qai_hub_compile_without_output_file(
    mock_console,
    mock_prepare_input_size,
    mock_compiler_class,
    ctx_obj,
    temp_model_file,
    monkeypatch,
):
    """Test QAI Hub compile without specifying output file."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_prepare_input_size.return_value = [1, 3, 224, 224]
    mock_compiler = Mock()
    mock_result = Mock()
    mock_result.device = "Samsung Galaxy S24"
    mock_compiler.compile.return_value = mock_result
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = [
        "qai-hub",
        "--model",
        str(temp_model_file),
        "--device",
        "Samsung Galaxy S24",
        "--size",
        "1,3,224,224",
        "--runtime",
        "tflite",
    ]
    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 0
    # Should generate output filename based on model file and runtime
    expected_output = temp_model_file.with_suffix(".tflite").as_posix()
    mock_console.print.assert_any_call(
        f"[yellow]No output file specified, using {expected_output}[/]"
    )


def test_qai_hub_compile_missing_required_args_fails(ctx_obj, monkeypatch):
    """Test that QAI Hub compile fails when required arguments are missing."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )
    result = runner.invoke(compile_cli, ["qai-hub"], obj=ctx_obj)
    assert result.exit_code == 2
    assert "Missing option" in result.output or "Usage:" in result.output


@patch("embedl_hub.core.compile.QAIHubCompiler")
@patch("embedl_hub.cli.utils.prepare_input_size")
@patch("embedl_hub.core.hub_logging.console")
def test_qai_hub_compile_handles_compile_error(
    mock_console,
    mock_prepare_input_size,
    mock_compiler_class,
    ctx_obj,
    temp_model_file,
    monkeypatch,
):
    """Test QAI Hub compile handles CompileError gracefully."""
    from embedl_hub.core.compile import CompileError

    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_prepare_input_size.return_value = [1, 3, 224, 224]
    mock_compiler = Mock()
    mock_compiler.compile.side_effect = CompileError("Compilation failed")
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = [
        "qai-hub",
        "--model",
        str(temp_model_file),
        "--device",
        "Samsung Galaxy S24",
        "--size",
        "1,3,224,224",
        "--runtime",
        "tflite",
    ]
    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 1
    mock_console.print.assert_called_with("[red]✗ Compilation failed[/]")


@patch("embedl_hub.core.compile.QAIHubCompiler")
@patch("embedl_hub.cli.utils.prepare_input_size")
@patch("embedl_hub.core.hub_logging.console")
def test_qai_hub_compile_handles_value_error(
    mock_console,
    mock_prepare_input_size,
    mock_compiler_class,
    ctx_obj,
    temp_model_file,
    monkeypatch,
):
    """Test QAI Hub compile handles ValueError gracefully."""
    monkeypatch.setattr("embedl_hub.cli.utils.assert_api_config", lambda: None)
    monkeypatch.setattr(
        "embedl_hub.core.context.require_initialized_ctx", lambda x: None
    )

    mock_prepare_input_size.return_value = [1, 3, 224, 224]
    mock_compiler = Mock()
    mock_compiler.compile.side_effect = ValueError("Invalid input size")
    mock_compiler_class.return_value = mock_compiler
    mock_compiler_class.validate_core_args = Mock()

    args = [
        "qai-hub",
        "--model",
        str(temp_model_file),
        "--device",
        "Samsung Galaxy S24",
        "--size",
        "1,3,224,224",
        "--runtime",
        "tflite",
    ]
    result = runner.invoke(compile_cli, args, obj=ctx_obj)

    assert result.exit_code == 1
    mock_console.print.assert_called_with("[red]✗ Invalid input size[/]")


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