# Copyright 2022-2024 ETSI SDG TeraFlowSDN (TFS) (https://tfs.etsi.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest
import logging
import pandas as pd
from unittest.mock import MagicMock, patch
from common.tools.kafka.Variables import KafkaTopic
from analytics.backend.service.Streamer import DaskStreamer
from .messages_analyzer import get_batch, get_input_kpi_list, get_output_kpi_list, get_thresholds, \
                                    get_windows_size, get_batch_size, get_agg_df
from analytics.backend.service.AnalyzerHandlers import aggregation_handler, threshold_handler

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(funcName)s -  %(levelname)s - %(message)s')

# --- "test_validate_kafka_topics" should be run before the functionality tests ---
def test_validate_kafka_topics():
    logger.debug(" >>> test_validate_kafka_topics: START <<< ")
    response = KafkaTopic.create_all_topics()
    assert isinstance(response, bool)

###########################
# Tests Implementation of Telemetry Backend
###########################

@pytest.fixture(autouse=True)
def log_all_methods(request):
    '''
    This fixture logs messages before and after each test function runs, indicating the start and end of the test.
    The autouse=True parameter ensures that this logging happens automatically for all tests in the module.
    '''
    logger.info(f" >>> Starting test: {request.node.name} >>> ")
    yield
    logger.info(f" <<< Finished test: {request.node.name} <<< ")

@pytest.fixture
def dask_streamer():
    with patch('analytics.backend.service.AnalyzerHelper.AnalyzerHelper.initialize_dask_client')    as mock_dask_client, \
         patch('analytics.backend.service.AnalyzerHelper.AnalyzerHelper.initialize_kafka_consumer') as mock_kafka_consumer, \
         patch('analytics.backend.service.AnalyzerHelper.AnalyzerHelper.initialize_kafka_producer') as mock_kafka_producer:
        
        mock_dask_client.return_value = (MagicMock(), MagicMock())
        mock_kafka_consumer.return_value = MagicMock()
        mock_kafka_producer.return_value = MagicMock()
        
        return DaskStreamer(
            key="test_key",
            input_kpis=get_input_kpi_list(),
            output_kpis=get_output_kpi_list(),
            thresholds=get_thresholds(),
            batch_size=get_batch_size(),
            window_size=get_windows_size(),
            n_workers=3,
            threads_per_worker=1
        )

def test_initialization(dask_streamer):
    """Test if the DaskStreamer initializes correctly."""
    assert dask_streamer.key == "test_key"
    assert dask_streamer.batch_size == get_batch_size()
    assert dask_streamer.window_size is None
    assert dask_streamer.n_workers == 3
    assert dask_streamer.consumer is not None
    assert dask_streamer.producer is not None
    assert dask_streamer.client is not None
    assert dask_streamer.cluster is not None


def test_run_stops_on_no_consumer(dask_streamer):
    """Test if the run method exits when the consumer is not initialized."""
    dask_streamer.consumer = None
    with patch('time.sleep', return_value=None):
        dask_streamer.run()
    assert not dask_streamer.running

def test_task_handler_selector_valid_handler(dask_streamer):
    """Test task handler selection with a valid handler."""
    with patch('analytics.backend.service.AnalyzerHandlers.AnalyzerHandlers.is_valid_handler', return_value=True), \
         patch.object(dask_streamer.client, 'submit', return_value=MagicMock()) as mock_submit, \
         patch.object(dask_streamer.client, 'status', 'running'):

        dask_streamer.task_handler_selector()
        mock_submit.assert_called_once()

def test_task_handler_selector_invalid_handler(dask_streamer):
    """Test task handler selection with an invalid handler."""
    with patch('analytics.backend.service.AnalyzerHandlers.AnalyzerHandlers.is_valid_handler', return_value=False):
        dask_streamer.task_handler_selector()
        assert dask_streamer.batch == []

def test_produce_result(dask_streamer):
    """Test if produce_result sends records to Kafka."""
    result = [{"kpi_id": "kpi1", "value": 100}]
    with patch('analytics.backend.service.AnalyzerHelper.AnalyzerHelper.delivery_report', return_value=None) as mock_delivery_report, \
         patch.object(dask_streamer.producer, 'produce') as mock_produce:
        dask_streamer.produce_result(result, "test_topic")
        mock_produce.assert_called_once_with(
            "test_topic",
            key="kpi1",
            value='{"kpi_id": "kpi1", "value": 100}',
            callback=mock_delivery_report
        )

def test_cleanup(dask_streamer):
    """Test the cleanup method."""
    with patch.object(dask_streamer.consumer, 'close')              as mock_consumer_close, \
         patch.object(dask_streamer.producer, 'flush')              as mock_producer_flush, \
         patch.object(dask_streamer.client,   'close')              as mock_client_close, \
         patch.object(dask_streamer.cluster,  'close', MagicMock()) as mock_cluster_close:
        
        # Mock the conditions required for the close calls
        dask_streamer.client.status = 'running'
        dask_streamer.cluster.close = MagicMock()
        
        dask_streamer.cleanup()

        mock_consumer_close.assert_called_once()
        mock_producer_flush.assert_called_once()
        mock_client_close.assert_called_once()
        dask_streamer.cluster.close.assert_called_once()

def test_run_with_valid_consumer(dask_streamer):
    """Test the run method with a valid Kafka consumer."""
    with patch.object(dask_streamer.consumer, 'poll') as mock_poll, \
         patch.object(dask_streamer, 'task_handler_selector') as mock_task_handler_selector:
        
        # Simulate valid messages without errors
        mock_message_1 = MagicMock()
        mock_message_1.value.return_value = b'{"kpi_id": "kpi1", "value": 100}'
        mock_message_1.error.return_value = None  # No error

        mock_message_2 = MagicMock()
        mock_message_2.value.return_value = b'{"kpi_id": "kpi2", "value": 200}'
        mock_message_2.error.return_value = None  # No error

        # Mock `poll` to return valid messages
        mock_poll.side_effect = [mock_message_1, mock_message_2]
        
        # Run the `run` method in a limited loop
        with patch('time.sleep', return_value=None):  # Mock `sleep` to avoid delays
            dask_streamer.running = True  # Ensure the streamer runs
            dask_streamer.batch_size = 2  # Set a small batch size for the test
            
            # Limit the loop by breaking it after one full processing cycle
            def stop_running_after_task_handler(*args, **kwargs):
                logger.info("Stopping the streamer after processing the first batch.")
                dask_streamer.running = False
            
            mock_task_handler_selector.side_effect = stop_running_after_task_handler
            
            # Execute the method
            dask_streamer.run()
        
        # Assertions
        assert len(dask_streamer.batch) == 0  # Batch should be cleared after processing
        mock_task_handler_selector.assert_called_once()  # Task handler should be called once
        mock_poll.assert_any_call(timeout=2.0)  # Poll should have been called


# add a test to check the working of aggregation_handler function and threshold_handler from AnalyzerHandlers.py
def test_aggregation_handler():

    # Create a sample batch
    batch = get_batch()
    input_kpi_list = get_input_kpi_list()
    output_kpi_list = get_output_kpi_list()
    thresholds = get_thresholds() 

    # Test aggregation_handler
    aggregated_df = aggregation_handler(
        "test_batch", "test_key", batch, input_kpi_list, output_kpi_list, thresholds
    )
    assert isinstance(aggregated_df, list)
    assert all(isinstance(item, dict) for item in aggregated_df)

# Test threshold_handler
def test_threshold_handler():
    # Create a sample aggregated DataFrame
    agg_df = get_agg_df()
    thresholds = get_thresholds()

    # Test threshold_handler
    result = threshold_handler("test_key", agg_df, thresholds["task_parameter"][0])
    assert isinstance(result, pd.DataFrame)
    assert result.shape == (1, 7)
