"""
Task management and execution system.

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

import json
import os
import random
import threading
import time

from typing import Any

from dependency_injector.wiring import Provide, inject

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.odbc.base import (
    BaseODBCConnectionConfig,
)
from data_exchange_agent.config.sections.connections.odbc.redshift import (
    RedshiftConnectionConfig,
)
from data_exchange_agent.config.sections.connections.odbc.sqlserver import (
    SQLServerConnectionConfig,
)
from data_exchange_agent.constants import config_keys, container_keys, task_keys
from data_exchange_agent.constants.config_defaults import (
    DEFAULT__APPLICATION__MAX_PARALLEL_TASKS,
    DEFAULT__APPLICATION__TASK_FETCH_INTERVAL,
)
from data_exchange_agent.constants.data_source_types import DataSourceType
from data_exchange_agent.constants.paths import (
    build_actual_results_folder_path,
)
from data_exchange_agent.data_sources import BaseDataSource, DataSourceRegistry
from data_exchange_agent.data_sources.database_engines import DatabaseEngine
from data_exchange_agent.data_sources.sf_connection import SnowflakeDataSource
from data_exchange_agent.data_sources.sql_command_type import SQLCommandType
from data_exchange_agent.data_sources.sql_parser import get_permitted_sql_command_type
from data_exchange_agent.interfaces.task_source_adapter import TaskSourceAdapter
from data_exchange_agent.providers.storage_provider import StorageProvider
from data_exchange_agent.tasks.data_validation import DataValidationHandler
from data_exchange_agent.tasks.lease_refresher import LeaseRefresher
from data_exchange_agent.utils.decorators import log_error
from data_exchange_agent.utils.sf_logger import SFLogger
from shared.enums.source_type import SourceType
from shared.telemetry import (
    FEATURE_DATA_MIGRATION,
    UNKNOWN_SOURCE,
    UNKNOWN_TARGET,
    MigrationTracker,
)
from snowflake.connector import errors as snowflake_errors


JITTER_MAX_SECONDS = 2.0

# Mapping from database engine to bulk utility data source type
# Only engines with available bulk utilities are included
ENGINE_TO_BULK_UTILITY: dict[str, str] = {
    DatabaseEngine.SQLSERVER.value: DataSourceType.BCP.value,
}


class TaskManager:
    """
    Manages task processing using persistent worker threads.

    This class handles task execution across multiple worker threads. Each worker
    directly pulls tasks from the task source (with capacity=1), processes them,
    and then pulls the next task. This ensures each worker only handles one task
    at a time and tasks are claimed at the source level.

    Attributes:
        debug_mode (bool): Whether to run in debug mode (single worker)
        max_parallel_tasks (int): Maximum number of parallel tasks
        task_fetch_interval (int): Interval in seconds to wait between task fetches when idle
        worker_threads (list[threading.Thread]): List of active worker threads
        stop_workers (bool): Flag to signal workers to stop
        handling_tasks (bool): Whether task handling is currently active
        agent_id (str): Agent ID for task source communication
        source_connections_config (dict): Database connection configurations
        target_connections_config (dict): Cloud storage connection configurations

    Args:
        max_parallel_tasks (int, optional): Maximum number of parallel tasks. Defaults to DEFAULT__APPLICATION__MAX_PARALLEL_TASKS.
        task_fetch_interval (int, optional): Interval in seconds between task fetches when idle.

    """

    @property
    def stop_workers(self) -> bool:
        """Get the stop workers flag."""
        return self._stop_workers

    @stop_workers.setter
    def stop_workers(self, value: bool) -> None:
        """Set the stop workers flag."""
        self._stop_workers = value

    @log_error
    @inject
    def __init__(
        self,
        max_parallel_tasks: int = DEFAULT__APPLICATION__MAX_PARALLEL_TASKS,
        task_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
        ],
        snowflake_datasource: SnowflakeDataSource = Provide[
            container_keys.SF_DATA_SOURCE
        ],
    ) -> None:
        """
        Initialize the TaskManager with specified configuration.

        Args:
            max_parallel_tasks (int): Maximum number of parallel tasks.
            task_fetch_interval (int): Interval in seconds to wait between task fetches when idle.
            logger (SFLogger): Logger instance for logging messages.
            program_config (ConfigManager): Program configuration.
            task_source_adapter (TaskSourceAdapter): Task source adapter instance.
            snowflake_datasource (SnowflakeDataSource): Shared Snowflake datasource instance.

        """
        self.logger = logger
        self.task_source_adapter = task_source_adapter
        self.snowflake_datasource = snowflake_datasource

        # Check if debug mode is enabled (for better debugging with breakpoints)
        self.debug_mode = os.getenv("DEBUG_SINGLE_WORKER") == "1"
        if self.debug_mode:
            self.logger.info("DEBUG MODE: Running with single worker for debugging")
            max_parallel_tasks = 1

        self.max_parallel_tasks = max_parallel_tasks
        self.task_fetch_interval = task_fetch_interval
        self.worker_threads: list[threading.Thread] = []
        self.stop_workers = False
        self.handling_tasks = False
        self.agent_id = program_config[config_keys.APPLICATION__AGENT_ID]

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

        # Lease refresher to keep task leases alive while workers are processing
        self.lease_refresher = LeaseRefresher()

    @log_error
    def handle_tasks(self) -> None:
        """
        Start task handling with persistent worker threads.

        Creates multiple worker threads, each of which directly pulls tasks
        from the task source and processes them. Each worker handles one
        task at a time, ensuring proper parallelism.
        """
        if self.handling_tasks:
            self.logger.warning(
                f"TaskManager already handling tasks in PID {os.getpid()} with agent ID '{self.agent_id}'."
            )
            return

        self.handling_tasks = True
        self.stop_workers = False
        self.logger.info(
            f"Starting task handling in PID {os.getpid()} with agent ID '{self.agent_id}' and {self.max_parallel_tasks} worker(s)."
        )

        # Start the lease refresher to keep task leases alive
        self.lease_refresher.start()

        # Start N persistent worker threads
        for i in range(self.max_parallel_tasks):
            worker_thread = threading.Thread(
                target=self._worker_loop, daemon=True, name=f"Worker-{i}"
            )
            self.worker_threads.append(worker_thread)
            worker_thread.start()
            self.logger.info(f"Started worker thread: {worker_thread.name}")

    def _worker_loop(self) -> None:
        """
        Worker loop - continuously pull and process tasks from the source.

        Each worker directly calls the task source to get a task (capacity=1),
        processes it, and then requests the next task. This ensures:
        - Each worker handles exactly one task at a time
        - Tasks are claimed at the source level (no local queue issues)
        - Workers only request new tasks when they're ready
        """
        worker_name = threading.current_thread().name
        self.logger.info(f"{worker_name} started and ready to process tasks.")

        while not self.stop_workers:
            try:
                # Pull task directly from source (source returns max 1 task)
                tasks = self.task_source_adapter.get_tasks()

                if tasks:
                    task = tasks[0]  # get_tasks returns a list, take the first one
                    task_id = str(task.get(task_keys.ID))
                    self.logger.info(f"{worker_name} claimed task: {task_id}")
                    self.process_task(task)
                else:
                    # No tasks available, wait before checking again
                    time.sleep(self._jittered_interval())

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

        # Close this worker thread's Snowflake connection (thread-local inside SnowflakeDataSource).
        try:
            self.snowflake_datasource.close_connection()
        except Exception:
            pass
        self.logger.info(f"{worker_name} stopped.")

    def stop(self) -> None:
        """Signal all workers to stop processing tasks."""
        self.logger.info("Stopping task handling...")
        self.stop_workers = True
        self.handling_tasks = False
        self.lease_refresher.stop()

    def _jittered_interval(self) -> float:
        """Return task_fetch_interval plus uniform random jitter to desynchronize workers."""
        return self.task_fetch_interval + random.uniform(0, JITTER_MAX_SECONDS)

    def process_task(self, task: dict[str, Any]) -> None:
        """
        Route a task to the appropriate handler based on its kind.

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

        """
        task_kind = task.get(task_keys.KIND)

        if task_kind == task_keys.DATA_VALIDATION:
            self._process_data_validation_task(task)
        else:
            self._process_data_movement_task(task)

    def _process_data_movement_task(self, task: dict[str, Any]) -> None:
        """Handle a data-movement (extraction + upload) task."""
        task_id = task.get(task_keys.ID)
        worker_name = threading.current_thread().name
        try:
            formatted_task = json.dumps(task, indent=4)
        except (TypeError, ValueError):
            formatted_task = str(task)
        self.logger.info(f"{worker_name} processing task '{task_id}': {formatted_task}")

        # Use MigrationTracker context manager for automatic telemetry
        with MigrationTracker(
            task_id=str(task_id),
            source_type=task.get(task_keys.ENGINE, UNKNOWN_SOURCE),
            target_type=task.get(task_keys.TARGET_TYPE, UNKNOWN_TARGET),
            feature=FEATURE_DATA_MIGRATION,
        ) as tracker:
            # Get the engine configuration
            engine = task.get(task_keys.ENGINE)
            if engine not in self.source_connections_config:
                message = (
                    f"Engine {engine} not found in source connections configuration."
                )
                self._handle_error(task, message)
                tracker.log_failure(error_msg=message)
                return

            # Create the data source
            data_source_name = None
            statement = task.get(task_keys.STATEMENT)
            target_type = task.get(task_keys.TARGET_TYPE)
            target_id = task.get(task_keys.TARGET_ID)

            try:
                # Handle S3 UNLOAD target type - DEA builds UNLOAD command
                if target_type == task_keys.TargetTypeValues.S3_UNLOAD:
                    if statement is None:
                        raise ValueError(
                            "Statement is required for S3_UNLOAD target type"
                        )
                    statement, data_source_name = self._prepare_unload_statement(
                        engine, statement, target_id
                    )
                else:
                    # Determine data source type based on is_bulk_operation flag and statement
                    is_bulk_operation = task.get(task_keys.IS_BULK_OPERATION, False)
                    data_source_name = self._resolve_data_source_type(
                        engine, is_bulk_operation, statement
                    )

                results_folder_path = build_actual_results_folder_path(task_id)
                base_file_name = f"task_{task_id}_result"
                suggested_batch_size = task.get(task_keys.SUGGESTED_BATCH_SIZE)
                database = task.get(task_keys.DATABASE)
                use_database_command = task.get(task_keys.USE_DATABASE_COMMAND)
                data_source: BaseDataSource = DataSourceRegistry.create(
                    data_source_name,
                    engine=engine,
                    statement=statement,
                    results_folder_path=results_folder_path,
                    base_file_name=base_file_name,
                    suggested_batch_size=suggested_batch_size,
                    database=database,
                    use_database_command=use_database_command,
                )
            except Exception as e:
                message = f"Failed creating the '{data_source_name}' data source."
                self._handle_error(task, message, e)
                tracker.log_failure(exception=e, error_msg=message)
                return

            # Export the data from the data source
            try:
                if not data_source.export_data():
                    raise Exception("Failed exporting data from the data source.")
            except Exception as e:
                message = f"Failed exporting data from the '{task[task_keys.ENGINE]}' engine using the '{data_source_name}' data source."
                self._handle_odbc_error(task, message, e)
                tracker.log_failure(exception=e, error_msg=message)
                return

            # Upload the data (skip for UNLOAD operations as data is already in S3)
            if data_source_name != DataSourceType.UNLOAD.value:
                try:
                    if target_type is None:
                        raise ValueError(
                            "target_type is required for upload operations"
                        )
                    if target_id is None:
                        raise ValueError("target_id is required for upload operations")
                    self.logger.info(
                        f"Uploading the results to the '{target_type}' cloud storage."
                    )
                    storage_provider = StorageProvider(
                        target_type, self.target_connections_config
                    )
                    storage_provider.upload_files(results_folder_path, target_id)
                    self.logger.info(
                        f"Uploaded the results to the '{target_type}' cloud storage."
                    )
                except Exception as e:
                    message = f"Failed uploading the results to the '{target_type}' cloud storage."
                    self._handle_cloud_storage_error(task, message, e)
                    tracker.log_failure(exception=e, error_msg=message)
                    return
            else:
                self.logger.info(
                    "Skipping upload phase for UNLOAD operation - data already written to S3."
                )

            # Calculate and set metrics for telemetry
            # For UNLOAD operations, metrics will be 0 since no local files are created
            rows_processed, bytes_processed = self._calculate_metrics(
                results_folder_path
            )
            tracker.set_metrics(rows_processed, bytes_processed)

            # Handle success
            self._handle_success(task)

    def _process_data_validation_task(self, task: dict[str, Any]) -> None:
        """Dispatch a data_validation task to the DataValidationHandler."""
        task_id = task.get(task_keys.ID)
        worker_name = threading.current_thread().name
        self.logger.info(f"{worker_name} processing data_validation task '{task_id}'")

        handler = DataValidationHandler(
            logger=self.logger,
            snowflake_datasource=self.snowflake_datasource,
            source_connections_config=self.source_connections_config,
        )
        result = handler.handle(task)

        if result.get(task_keys.STATUS) == "error":
            self._handle_error(task, result.get(task_keys.DETAILS, "Validation failed"))
        else:
            self._handle_success(task)

    def _prepare_unload_statement(
        self, engine: str, select_query: str, target_id: str | None
    ) -> tuple[str, str]:
        """
        Prepare UNLOAD statement using DEA configuration.

        The orchestrator sends a SELECT query and the relative path (in target_id).
        The DEA builds the full UNLOAD command using its configured S3 bucket and IAM role.

        Args:
            engine: Database engine (must be 'redshift')
            select_query: The SELECT query to wrap in UNLOAD
            target_id: Relative path for UNLOAD output (e.g., "1/2/partition_0/")

        Returns:
            Tuple of (unload_statement, data_source_type)

        Raises:
            ValueError: If engine is not Redshift or UNLOAD is not configured

        """
        if SourceType.from_string(engine) != SourceType.REDSHIFT:
            raise ValueError(
                f"S3_UNLOAD target type is only supported for Redshift, got '{engine}'"
            )

        if not target_id:
            raise ValueError(
                "S3_UNLOAD target type requires 'target_id' with the relative S3 path"
            )

        # Get Redshift connection config and its UNLOAD config
        connection_config = self.source_connections_config.get(engine)
        if not isinstance(connection_config, RedshiftConnectionConfig):
            raise ValueError(
                f"Expected RedshiftConnectionConfig for engine '{engine}', got {type(connection_config).__name__}"
            )

        if not connection_config.unload_config:
            raise ValueError(
                "UNLOAD extraction strategy requires UNLOAD configuration in the DEA config. "
                "Add 'unload_s3_bucket', 'unload_iam_role_arn', and optionally 'unload_s3_prefix' "
                "to the [connections.redshift] section."
            )

        unload_config = connection_config.unload_config

        # Build full S3 path and UNLOAD query
        s3_path = unload_config.build_s3_path(target_id)
        unload_statement = unload_config.build_unload_query(select_query, s3_path)

        self.logger.info(
            f"Built UNLOAD statement:\n"
            f"  S3 path: {s3_path}\n"
            f"  Format: {unload_config.file_format}\n"
            f"  Max file size: {unload_config.max_file_size_mb} MB"
        )

        return unload_statement, DataSourceType.UNLOAD.value

    def _resolve_data_source_type(
        self, engine: str, is_bulk_operation: bool, statement: str | None = None
    ) -> str:
        """
        Resolve the data source type to use based on bulk flag and connection configuration.

        If is_bulk_operation is True and the connection has use_bcp=True,
        use BCP data source. Otherwise, use ODBC.

        For Redshift UNLOAD commands, use the UNLOAD data source type.

        Args:
            engine: Database engine (e.g., 'sqlserver', 'postgresql', 'redshift')
            is_bulk_operation: Whether bulk extraction is requested
            statement: SQL statement to execute (optional, used for UNLOAD detection)

        Returns:
            The data source type to use (e.g., 'odbc', 'bcp', 'unload')

        """
        # Check if this is a Redshift UNLOAD operation
        if statement:
            sql_command_type = get_permitted_sql_command_type(statement)
            if sql_command_type == SQLCommandType.UNLOAD:
                if engine.lower() == "redshift":
                    self.logger.info("Using UNLOAD data source for Redshift")
                    return DataSourceType.UNLOAD.value
                else:
                    self.logger.warning(
                        f"UNLOAD command detected but engine is '{engine}', not 'redshift'. "
                        f"UNLOAD is only supported for Redshift. Falling back to ODBC."
                    )

        # Default to ODBC for non-bulk operations
        if not is_bulk_operation:
            return DataSourceType.ODBC.value

        # Check if a bulk utility is available for this engine
        bulk_utility_type = ENGINE_TO_BULK_UTILITY.get(engine)
        if not bulk_utility_type:
            self.logger.debug(
                f"No bulk utility available for engine '{engine}', using ODBC"
            )
            return DataSourceType.ODBC.value

        # Check if BCP is enabled in the connection configuration
        connection_config = self.source_connections_config.get(engine)
        if (
            not isinstance(connection_config, SQLServerConnectionConfig)
            or not connection_config.use_bcp
        ):
            self.logger.debug(f"BCP not enabled for engine '{engine}', using ODBC")
            return DataSourceType.ODBC.value

        self.logger.info(f"Using BCP for engine '{engine}'")
        return bulk_utility_type

    def _handle_success(self, task: dict[str, Any]) -> None:
        task_id = task[task_keys.ID]
        worker_name = threading.current_thread().name
        try:
            self.task_source_adapter.complete_task(task_id)
            self.logger.info(f"Worker {worker_name} completed task '{task_id}'")
        except Exception as e:
            self._handle_error(
                task, f"Failed handling completion of task '{task_id}'.", e
            )

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

        # Build detailed error message if exception is provided
        full_message = (
            base_message
            if exception is None
            else self._build_detailed_error_message(base_message, exception)
        )
        self.logger.error(full_message, exception=exception)
        try:
            self.task_source_adapter.fail_task(task_id, full_message)
        except Exception as e:
            self.logger.error(
                f"Failed handling failure of task '{task_id}'.", exception=e
            )

    def _build_detailed_error_message(
        self, base_message: str, exception: Exception
    ) -> str:
        """
        Build a detailed error message including exception type and message.

        Args:
            base_message: The base error message to prepend
            exception: The exception to extract details from

        Returns:
            A formatted error message with complete exception details

        """
        exception_type = type(exception).__name__
        exception_message = str(exception)

        return f"{base_message}\n\nException Type: {exception_type}\nException Message: {exception_message}"

    def _handle_odbc_error(
        self, task: dict[str, Any], message: str, exception: Exception | None = None
    ) -> None:
        import pyodbc

        if isinstance(exception, pyodbc.Error):
            exception_message = str(exception)
            if (
                "Connection refused" in exception_message
                or "Login timeout" in exception_message
            ):
                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, exception)

    def _calculate_metrics(self, results_folder_path: str) -> tuple[int, int]:
        """
        Calculate rows and bytes processed from result files.

        Args:
            results_folder_path: Path to the results folder

        Returns:
            Tuple of (rows_processed, bytes_processed)

        """
        rows_processed = 0
        bytes_processed = 0

        try:
            import glob

            # Find all result files
            result_files = glob.glob(os.path.join(results_folder_path, "*.parquet"))

            for file_path in result_files:
                try:
                    # Get file size
                    bytes_processed += os.path.getsize(file_path)

                    # Try to count rows from parquet file
                    try:
                        import pyarrow.parquet as pq

                        parquet_file = pq.ParquetFile(file_path)
                        rows_processed += parquet_file.metadata.num_rows
                    except Exception:
                        # If we can't read parquet metadata, estimate based on file size
                        # Average ~1KB per row is a rough estimate
                        rows_processed += os.path.getsize(file_path) // 1024
                except Exception as e:
                    self.logger.debug(f"Error calculating metrics for {file_path}: {e}")

        except Exception as e:
            self.logger.debug(f"Error calculating metrics: {e}")

        return rows_processed, bytes_processed
