# 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 logging
import time
import json
from confluent_kafka import KafkaException, KafkaError
# import pandas as pd
from common.tools.kafka.Variables import KafkaTopic
from .AnalyzerHandlers import AnalyzerHandlers, aggregation_handler
from .AnalyzerHelper import AnalyzerHelper


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


class DaskStreamer:
    def __init__(self, key, input_kpis, output_kpis, thresholds, batch_size=5, 
                 window_size=None, n_workers=5, threads_per_worker=2):
        self.key         = key
        self.input_kpis  = input_kpis
        self.output_kpis = output_kpis
        self.thresholds  = thresholds
        self.window_size = window_size
        self.batch_size  = batch_size
        self.n_workers   = n_workers
        self.threads_per_worker = threads_per_worker
        self.running     = True
        self.batch       = []

        # Initialize Kafka and Dask components
        self.client, self.cluster = AnalyzerHelper.initialize_dask_client(n_workers, threads_per_worker)
        self.consumer = AnalyzerHelper.initialize_kafka_consumer()
        self.producer = AnalyzerHelper.initialize_kafka_producer()
        logger.info("Dask Streamer initialized.")

    def run(self):
        """Main method to start the DaskStreamer."""
        try:
            logger.info("Starting Dask Streamer")
            last_batch_time = time.time()
            while True:
                if not self.consumer:
                    logger.warning("Kafka consumer is not initialized or stopped. Exiting loop.")
                    break
                if not self.running:
                    logger.warning("Dask Streamer is not running. Exiting loop.")
                    break
                message = self.consumer.poll(timeout=2.0)
                if message is None:
                    # logger.info("No new messages received.")
                    continue
                if message.error():
                    if message.error().code() == KafkaError._PARTITION_EOF:
                        logger.warning(f"Consumer reached end of topic {message.topic()}/{message.partition()}")
                    elif message.error():
                        raise KafkaException(message.error())
                else:
                    try:
                        value = json.loads(message.value())
                    except json.JSONDecodeError:
                        logger.error(f"Failed to decode message: {message.value()}")
                        continue
                    self.batch.append(value)
                    # logger.info(f"Received message: {value}")

                # Window size has a priority over batch size 
                if self.window_size is None:
                    if len(self.batch) >= self.batch_size:  # If batch size is not provided, process continue with default batch size
                        logger.info(f"Processing based on batch size {self.batch_size}.")
                        self.task_handler_selector()
                        self.batch = []
                else:
                    # Process based on window size
                    current_time = time.time()
                    if (current_time - last_batch_time) >= self.window_size and self.batch:
                        logger.info(f"Processing based on window size {self.window_size}.")
                        self.task_handler_selector()
                        self.batch = []
                        last_batch_time = current_time

        except Exception as e:
            logger.exception(f"Error in Dask streaming process: {e}")
        finally:
            logger.info(">>> Exiting Dask Streamer...")
            self.cleanup()
            logger.info(">>> Dask Streamer Cleanup Completed.")

    def task_handler_selector(self):
        """Select the task handler based on the task type."""
        if AnalyzerHandlers.is_valid_handler(self.thresholds["task_type"]):
            if self.client.status == 'running':
                future = self.client.submit(aggregation_handler, "batch size", self.key,
                                                self.batch, self.input_kpis, self.output_kpis, self.thresholds)
                future.add_done_callback(lambda fut: self.produce_result(fut.result(), KafkaTopic.ALARMS.value))
            else:
                logger.warning("Dask client is not running. Skipping processing.")
        else:
            logger.warning(f"Unknown task type: {self.thresholds['task_type']}. Skipping processing.")

    def produce_result(self, result, destination_topic):
        """Produce results to the Kafka topic."""
        for record in result:
            try:
                self.producer.produce(
                    destination_topic,
                    key=str(record.get('kpi_id', '')),
                    value=json.dumps(record),
                    callback=AnalyzerHelper.delivery_report
                )
            except KafkaException as e:
                logger.error(f"Failed to produce message: {e}")
        self.producer.flush()
        logger.info(f"Produced {len(result)} aggregated records to '{destination_topic}'.")

    def cleanup(self):
        """Clean up Kafka and Dask resources."""
        logger.info("Shutting down resources...")
        self.running = False
        if self.consumer:
            try:
                self.consumer.close()
                logger.info("Kafka consumer closed.")
            except Exception as e:
                logger.error(f"Error closing Kafka consumer: {e}")

        if self.producer:
            try:
                self.producer.flush()
                logger.info("Kafka producer flushed and closed.")
            except Exception as e:
                logger.error(f"Error closing Kafka producer: {e}")

        if self.client and hasattr(self.client, 'status') and self.client.status == 'running':
            try:
                self.client.close()
                logger.info("Dask client closed.")
            except Exception as e:
                logger.error(f"Error closing Dask client: {e}")

        if self.cluster and hasattr(self.cluster, 'close'):
            try:
                self.cluster.close(timeout=5)
                logger.info("Dask cluster closed.")
            except Exception as e:
                logger.error(f"May be timeout. Error closing Dask cluster: {e}")
