"""
Task management and execution system.

This module provides the TaskManager class which orchestrates task execution
across multiple worker threads, handles task lifecycle management, API
communication, and integrates with various data sources and uploaders.
"""

import json
import os
import threading
import time

from concurrent.futures import ThreadPoolExecutor
from typing import Any

from dependency_injector.wiring import Provide, inject
from py4j.protocol import Py4JJavaError

from data_exchange_agent.config.manager import ConfigManager
from data_exchange_agent.config.sections.connections.cloud_storages.base import BaseCloudStorageConnectionConfig
from data_exchange_agent.config.sections.connections.jdbc.base import BaseJDBCConnectionConfig
from data_exchange_agent.constants import config_keys, container_keys, task_keys
from data_exchange_agent.constants.config_defaults import (
    DEFAULT__APPLICATION__TASK_FETCH_INTERVAL,
    DEFAULT__APPLICATION__WORKERS,
)
from data_exchange_agent.constants.paths import (
    build_actual_results_folder_path,
)
from data_exchange_agent.data_sources.extractor import DataSourceExtractor
from data_exchange_agent.data_sources.pyspark import PySparkDataSource
from data_exchange_agent.interfaces.data_source import DataSourceInterface
from data_exchange_agent.interfaces.task_queue import TaskQueueInterface
from data_exchange_agent.interfaces.task_source_adapter import TaskSourceAdapter
from data_exchange_agent.providers.storageProvider import StorageProvider
from data_exchange_agent.queues.sqlite_task_queue import SQLiteTaskQueue
from data_exchange_agent.utils.decorators import log_error
from data_exchange_agent.utils.sf_logger import SFLogger
from snowflake.connector import errors as snowflake_errors


class TaskManager:
    """
    Manages asynchronous task processing using a thread pool and task queue.

    This class handles fetching, queueing and processing of data extraction tasks.
    It maintains a thread pool for parallel task execution and uses a thread-safe
    task queue.

    Attributes:
        executor (ThreadPoolExecutor): Thread pool for executing tasks
        task_queue (TaskQueueInterface): Thread-safe queue of tasks to process
        stop_queue (bool): Flag to stop task processing
        task_source_adapter (TaskSourceAdapter): Adapter for task source interactions
        tasks_fetch_interval (int): Seconds between task fetch attempts
        handling_tasks (bool): Whether tasks are currently being handled
        connection_toml (dict): Database connection configurations

    Args:
        workers (int, optional): Number of worker threads. Defaults to DEFAULT_WORKERS_COUNT.
        tasks_fetch_interval (int, optional): Task fetch interval in seconds. Defaults to DEFAULT_TASKS_FETCH_INTERVAL.

    """

    @property
    def stop_queue(self) -> bool:
        """
        Get the stop queue flag.

        Returns:
            bool: True if queue processing should stop, False otherwise

        """
        return self._stop_queue

    @stop_queue.setter
    def stop_queue(self, value: bool) -> None:
        """
        Set the stop queue flag.

        Args:
            value (bool): True to stop queue processing, False to continue

        """
        self._stop_queue = value

    @log_error
    @inject
    def __init__(
        self,
        workers: int = DEFAULT__APPLICATION__WORKERS,
        tasks_fetch_interval: int = DEFAULT__APPLICATION__TASK_FETCH_INTERVAL,
        logger: SFLogger = Provide[container_keys.SF_LOGGER],
        program_config: ConfigManager = Provide[container_keys.PROGRAM_CONFIG],
        task_source_adapter: TaskSourceAdapter = Provide[container_keys.TASK_SOURCE_ADAPTER],
    ) -> None:
        """
        Initialize the TaskManager with specified configuration.

        Sets up the task execution environment including thread pool executor,
        task queue, task source adapter, and loads configuration from TOML file.
        Initializes all necessary components for managing and processing tasks.

        Args:
            workers (int): Number of worker threads for concurrent task execution.
                         Defaults to DEFAULT_WORKERS_COUNT.
            tasks_fetch_interval (int): Interval in seconds between API task fetches.
                                      Defaults to DEFAULT_TASKS_FETCH_INTERVAL seconds.
            logger (SFLogger): Logger instance for logging messages.
                             Defaults to injected sf_logger.
            program_config (ConfigManager): Program configuration
            task_source_adapter (TaskSourceAdapter): Task source adapter instance.
                             Defaults to injected task_source_adapter.

        Raises:
            Exception: If the configuration TOML file cannot be loaded.
            custom_exceptions.ConfigurationError: If something was wrong with the configuration TOML file.

        """
        self.logger = logger
        self.executor = ThreadPoolExecutor(max_workers=workers)
        self.task_queue: TaskQueueInterface = SQLiteTaskQueue()
        self.stop_queue = False
        self.task_source_adapter = task_source_adapter
        self.tasks_fetch_interval = tasks_fetch_interval
        self.handling_tasks = False

        self.source_connections_config: dict[str, BaseJDBCConnectionConfig] = program_config[
            config_keys.CONNECTIONS__SOURCE
        ]
        self.target_connections_config: dict[str, BaseCloudStorageConnectionConfig] = program_config[
            config_keys.CONNECTIONS__TARGET
        ]

        self.sources = {"jdbc_pyspark": PySparkDataSource}

    @log_error
    def add_task(self, task: dict[str, Any]) -> None:
        """
        Add a single task to the task queue.

        Args:
            task (dict[str, any]): Task configuration dictionary

        """
        self.task_queue.add_task(task)

    @log_error
    def get_tasks(self) -> None:
        """Fetch tasks from API and add them to the task queue."""
        tasks = self.task_source_adapter.get_tasks()
        for task in tasks:
            self.task_queue.add_task(task)

    @log_error
    def get_tasks_count(self) -> int:
        """
        Get current number of tasks in queue.

        Returns:
            int: Number of tasks in queue

        """
        return self.task_queue.get_queue_size()

    @log_error
    def get_deque_id(self) -> int:
        """
        Get memory ID of task queue.

        Returns:
            int: Memory ID of task queue object

        """
        return id(self.task_queue)

    @log_error
    def get_completed_count(self) -> int:
        """
        Get the number of completed tasks.

        Returns:
            Number of completed tasks

        """
        return self.task_queue.get_completed_count()

    @log_error
    def handle_tasks(self) -> None:
        """
        Start task handling in a background thread.

        Creates a daemon thread to process tasks if not already running.
        """
        if self.handling_tasks:
            self.logger.info(f"*** TaskManager already handling tasks in PID: {os.getpid()} ***")
            return
        self.handling_tasks = True
        task_thread = threading.Thread(target=self.task_loop, daemon=True)
        task_thread.start()

    def task_loop(self) -> None:
        """
        Process tasks continuously in a loop.

        Continuously fetches and processes tasks from the queue until stopped.
        Handles task retrieval, execution and error handling.
        """
        while True:
            try:
                if self.stop_queue:
                    self.handling_tasks = False
                    self.stop_queue = False
                    break

                self.get_tasks()

                while True:
                    task = None
                    task = self.task_queue.get_task()

                    if task:
                        formatted_task = json.dumps(task, indent=4)
                        self.logger.info(f"🚀 Processing task '{task[task_keys.ID]}': {formatted_task}")
                        self.executor.submit(
                            self.process_task,
                            task,
                        )
                        time.sleep(0.5)
                    else:
                        break

            except Exception as e:
                self.logger.error(f"Error in task_loop: {str(e)}", exception=e)
            finally:
                time.sleep(self.tasks_fetch_interval)

    def process_task(self, task: dict[str, any]) -> None:
        """
        Process a single data extraction task.

        Creates appropriate data source, extracts data and saves to parquet.
        Updates task status on completion.

        Args:
            task (dict[str, any]): Task configuration dictionary

        """
        # Get the engine configuration
        engine = task[task_keys.ENGINE]
        if engine in self.source_connections_config:
            engine_config: BaseJDBCConnectionConfig = self.source_connections_config[engine]
        else:
            message = f"Engine {engine} not found in source connections configuration."
            self._handle_error(task, message)
            return

        # Create the data source
        try:
            data_source_class = self.sources[task[task_keys.SOURCE_TYPE]]
            data_source: DataSourceInterface = data_source_class(engine_config)
            data_source.create_connection()
        except Exception as e:
            message = f"Failed creating a connection to the '{task[task_keys.ENGINE]}' data source."
            self._handle_error(task, message, e)
            return

        # Extract the data
        try:
            data_source_extractor = DataSourceExtractor(
                statement=task[task_keys.STATEMENT],
                connection=data_source,
            )
            results_folder_path = build_actual_results_folder_path(task[task_keys.ID])
            data_source_extractor.extract_data(results_folder_path)
        except Exception as e:
            message = (
                f"Failed executing the query '{task[task_keys.STATEMENT]}' "
                f"in the '{task[task_keys.ENGINE]}' data source."
            )
            self._handle_jdbc_error(task, message, e)
            return

        # Upload the data
        try:
            storage_provider = StorageProvider(task[task_keys.UPLOAD_TYPE], self.target_connections_config)
            storage_provider.upload_files(results_folder_path, task[task_keys.UPLOAD_PATH])
        except Exception as e:
            message = f"Failed uploading the results to the '{task[task_keys.UPLOAD_TYPE]}' cloud storage."
            self._handle_cloud_storage_error(task, message, e)
            return

        # Handle success
        self._handle_success(task)

    def _handle_success(self, task: dict[str, any]) -> None:
        task_id = task[task_keys.ID]
        self.task_source_adapter.complete_task(task_id)
        self.task_queue.complete_task(task_id)

    def _handle_error(self, task: dict[str, any], message: str, exception: Exception | None = None) -> None:
        task_id = task[task_keys.ID]
        message = f"Error processing the task '{task_id}'. {message}"
        self.logger.error(message, exception=exception)
        self.task_source_adapter.fail_task(task_id, message)
        self.task_queue.fail_task(task_id, message)

    def _handle_jdbc_error(self, task: dict[str, any], message: str, exception: Exception | None = None) -> None:
        if isinstance(exception, Py4JJavaError):
            exception_message = str(exception)
            if (
                "com.microsoft.sqlserver.jdbc.SQLServerException" in exception_message  # SQL Server specific exception
                and "Connection refused" in exception_message  # Connection refused exception
            ):
                message += (
                    "\nVerify the connection properties. "
                    "Make sure that an instance of SQL Server is running on the host "
                    "and accepting TCP/IP connections at the port. "
                    "Make sure that TCP connections to the port are not blocked by a firewall."
                )
                self._handle_error(task, message)
                return

        message += "\nPlease check the database connection and the query."
        self._handle_error(task, message, exception)

    def _handle_cloud_storage_error(
        self, task: dict[str, any], message: str, exception: Exception | None = None
    ) -> None:
        if isinstance(exception, snowflake_errors.DatabaseError | snowflake_errors.HttpError):
            message += f" {str(exception)}"
            self._handle_error(task, message)
        else:
            self._handle_error(task, message, exception)
