# Copyright (C) 2025 Embedl AB

"""
Tests for tflite_utils module.
"""

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

import pytest
from tensorflow.lite.profiling.proto import profiling_info_pb2

from embedl_hub.core.utils.tflite_utils import (
    _count_layers_by_unit,
    parse_tflite_profiling_artifacts,
)


@patch.object(
    profiling_info_pb2,
    "BenchmarkProfilingData",
)
@patch('embedl_hub.core.utils.tflite_utils.json_format')
@patch('embedl_hub.core.utils.tflite_utils._parse_memory_from_log_file')
def test_parse_valid_profiling_proto(
    mock_parse_memory, mock_json_format, mock_profiling_pb2
):
    """Test parsing a valid profiling proto with single subgraph and layer."""
    # Create mock profiling data
    mock_profile_info = {
        'runtimeProfile': {
            'subgraphProfiles': [
                {
                    'perOpProfiles': [
                        {
                            'name': 'conv2d_layer_1',
                            'nodeType': 'CONV_2D',
                            'inferenceMicroseconds': {'avg': 1500.0},  # 1.5 ms
                        }
                    ]
                }
            ]
        }
    }

    # Set up mocks
    mock_proto = MagicMock()
    mock_profiling_pb2.BenchmarkProfilingData.return_value = mock_proto
    mock_json_format.MessageToDict.return_value = mock_profile_info
    mock_parse_memory.return_value = 128.5  # Mock memory usage

    with tempfile.NamedTemporaryFile(delete=False) as temp_proto_file:
        proto_path = Path(temp_proto_file.name)
        with tempfile.NamedTemporaryFile(delete=False) as temp_log_file:
            log_path = Path(temp_log_file.name)
            summary_dict, execution_detail = parse_tflite_profiling_artifacts(
                proto_path, log_path
            )

        # Verify summary dictionary structure
        assert 'mean_ms' in summary_dict
        assert 'peak_memory_usage_mb' in summary_dict
        assert 'layer_times' in summary_dict
        assert 'layers' in summary_dict
        assert 'layers_by_unit' in summary_dict

        # Verify summary values
        assert summary_dict['mean_ms'] == 1.5
        assert summary_dict['peak_memory_usage_mb'] == 128.5
        assert len(summary_dict['layer_times']) == 1
        assert summary_dict['layer_times'][0] == (
            1.5,
            'conv2d_layer_1',
            'CONV_2D',
        )

        # Verify execution details
        assert len(execution_detail) == 1
        layer_record = execution_detail[0]
        assert layer_record['name'] == 'conv2d_layer_1'
        assert layer_record['type'] == 'CONV_2D'
        assert layer_record['compute_unit'] == 'CPU'
        assert layer_record['execution_time'] == 1.5
        assert layer_record['execution_cycles'] == 0

        # Verify layers by unit
        assert summary_dict['layers_by_unit']['CPU'] == 1
        assert summary_dict['layers_by_unit']['GPU'] == 0
        assert summary_dict['layers_by_unit']['NPU'] == 0

    # Cleanup
    proto_path.unlink()
    log_path.unlink()


@patch.object(
    profiling_info_pb2,
    "BenchmarkProfilingData",
)
@patch('embedl_hub.core.utils.tflite_utils.json_format')
@patch('embedl_hub.core.utils.tflite_utils._parse_memory_from_log_file')
def test_parse_multiple_layers_and_subgraphs(
    mock_parse_memory, mock_json_format, mock_profiling_pb2
):
    """Test parsing proto with multiple subgraphs and layers."""
    mock_profile_info = {
        'runtimeProfile': {
            'subgraphProfiles': [
                {
                    'perOpProfiles': [
                        {
                            'name': 'input_layer',
                            'nodeType': 'INPUT',
                            'inferenceMicroseconds': {'avg': 100.0},  # 0.1 ms
                        },
                        {
                            'name': 'conv2d_layer_1',
                            'nodeType': 'CONV_2D',
                            'inferenceMicroseconds': {'avg': 2000.0},  # 2.0 ms
                        },
                    ]
                },
                {
                    'perOpProfiles': [
                        {
                            'name': 'relu_layer_1',
                            'nodeType': 'RELU',
                            'inferenceMicroseconds': {'avg': 500.0},  # 0.5 ms
                        }
                    ]
                },
            ]
        }
    }

    # Set up mocks
    mock_proto = MagicMock()
    mock_profiling_pb2.BenchmarkProfilingData.return_value = mock_proto
    mock_json_format.MessageToDict.return_value = mock_profile_info
    mock_parse_memory.return_value = 256.0  # Mock memory usage

    with tempfile.NamedTemporaryFile(delete=False) as temp_proto_file:
        proto_path = Path(temp_proto_file.name)
        with tempfile.NamedTemporaryFile(delete=False) as temp_log_file:
            log_path = Path(temp_log_file.name)
            summary_dict, execution_detail = parse_tflite_profiling_artifacts(
                proto_path, log_path
            )

        # Verify total execution time (sum of all layers)
        expected_total = (100.0 + 2000.0 + 500.0) / 1000.0  # Convert to ms
        assert summary_dict['mean_ms'] == expected_total

        # Verify we have 3 layers total
        assert len(execution_detail) == 3
        assert len(summary_dict['layer_times']) == 3

        # Verify top 5 layers are sorted by execution time (descending)
        layer_times = summary_dict['layer_times']
        assert layer_times[0][0] == 2.0  # conv2d_layer_1 is slowest
        assert layer_times[1][0] == 0.5  # relu_layer_1 is second
        assert layer_times[2][0] == 0.1  # input_layer is fastest

        # Verify layers_by_unit counts all 3 layers
        assert summary_dict['layers_by_unit']['CPU'] == 3

    proto_path.unlink()
    log_path.unlink()


@patch.object(
    profiling_info_pb2,
    "BenchmarkProfilingData",
)
@patch('embedl_hub.core.utils.tflite_utils.json_format')
@patch('embedl_hub.core.utils.tflite_utils._parse_memory_from_log_file')
def test_parse_top_5_layers_limit(
    mock_parse_memory, mock_json_format, mock_profiling_pb2
):
    """Test that only top 5 slowest layers are included in layer_times."""
    # Create 7 layers to test the top 5 limit
    layers = []
    for i in range(7):
        layers.append(
            {
                'name': f'layer_{i}',
                'nodeType': 'CONV_2D',
                'inferenceMicroseconds': {
                    'avg': float((i + 1) * 1000)
                },  # 1ms, 2ms, ..., 7ms
            }
        )

    mock_profile_info = {
        'runtimeProfile': {'subgraphProfiles': [{'perOpProfiles': layers}]}
    }

    # Set up mocks
    mock_proto = MagicMock()
    mock_profiling_pb2.BenchmarkProfilingData.return_value = mock_proto
    mock_json_format.MessageToDict.return_value = mock_profile_info
    mock_parse_memory.return_value = 512.0  # Mock memory usage

    with tempfile.NamedTemporaryFile(delete=False) as temp_proto_file:
        proto_path = Path(temp_proto_file.name)
        with tempfile.NamedTemporaryFile(delete=False) as temp_log_file:
            log_path = Path(temp_log_file.name)
            summary_dict, execution_detail = parse_tflite_profiling_artifacts(
                proto_path, log_path
            )

        # Verify only top 5 layers in layer_times
        assert len(summary_dict['layer_times']) == 5

        # Verify they are the 5 slowest (7ms, 6ms, 5ms, 4ms, 3ms)
        expected_times = [7.0, 6.0, 5.0, 4.0, 3.0]
        actual_times = [lt[0] for lt in summary_dict['layer_times']]
        assert actual_times == expected_times

        # But all 7 layers should be in execution_detail
        assert len(execution_detail) == 7

    proto_path.unlink()
    log_path.unlink()


def test_count_layers_by_unit_multiple_subgraphs():
    """Test counting layers across multiple subgraphs."""
    execution_detail = {
        'runtimeProfile': {
            'subgraphProfiles': [
                {
                    'perOpProfiles': [
                        {'name': 'layer1', 'nodeType': 'CONV_2D'},
                        {'name': 'layer2', 'nodeType': 'RELU'},
                    ]
                },
                {'perOpProfiles': [{'name': 'layer3', 'nodeType': 'POOL'}]},
            ]
        }
    }

    result = _count_layers_by_unit(execution_detail)

    assert result['CPU'] == 3  # Total of 3 layers
    assert result['GPU'] == 0
    assert result['NPU'] == 0


@patch.object(
    profiling_info_pb2,
    "BenchmarkProfilingData",
)
@patch('embedl_hub.core.utils.tflite_utils.json_format')
@patch('embedl_hub.core.utils.tflite_utils._parse_memory_from_log_file')
def test_parse_no_memory_info(
    mock_parse_memory, mock_json_format, mock_profiling_pb2
):
    """Test parsing when memory info is not available (returns None)."""
    mock_profile_info = {
        'runtimeProfile': {
            'subgraphProfiles': [
                {
                    'perOpProfiles': [
                        {
                            'name': 'test_layer',
                            'nodeType': 'TEST',
                            'inferenceMicroseconds': {'avg': 1000.0},  # 1.0 ms
                        }
                    ]
                }
            ]
        }
    }

    # Set up mocks
    mock_proto = MagicMock()
    mock_profiling_pb2.BenchmarkProfilingData.return_value = mock_proto
    mock_json_format.MessageToDict.return_value = mock_profile_info
    mock_parse_memory.return_value = None  # No memory info available

    with tempfile.NamedTemporaryFile(delete=False) as temp_proto_file:
        proto_path = Path(temp_proto_file.name)
        with tempfile.NamedTemporaryFile(delete=False) as temp_log_file:
            log_path = Path(temp_log_file.name)
            summary_dict, _ = parse_tflite_profiling_artifacts(
                proto_path, log_path
            )

        # Verify memory usage defaults to 0.0 when None is returned
        assert summary_dict['peak_memory_usage_mb'] == 0.0
        assert summary_dict['mean_ms'] == 1.0

    proto_path.unlink()
    log_path.unlink()


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