"""
Snowflake stored procedure task source adapter.

This module provides the SnowflakeStoredProcedureTaskSourceAdapter class for managing
communication with Snowflake stored procedures, including task retrieval and status updates.
"""

import json
import logging

from dependency_injector.wiring import Provide, inject

from data_exchange_agent import custom_exceptions
from data_exchange_agent.config import ConfigManager
from data_exchange_agent.constants import config_keys, container_keys, task_keys
from data_exchange_agent.data_sources.sf_connection import SnowflakeDataSource
from data_exchange_agent.interfaces.task_source_adapter import TaskSourceAdapter
from data_exchange_agent.snowflake_metadata_locations import (
    qualified_data_migration_metadata_schema,
)
from data_exchange_agent.tasks.task import Task
from snowflake.connector.errors import ProgrammingError


logger = logging.getLogger(__name__)

TRANSIENT_SF_TRANSACTION_ERROR_CODE = 90232
# Covers partial deployments where PULL_TASKS or its sub-procedures
# (PULL_SINGLE_TASK / PULL_MANY_TASKS) don't exist yet.
SF_UNKNOWN_FUNCTION_ERROR_CODE = 2141


class SnowflakeStoredProcedureTaskSourceAdapter(TaskSourceAdapter):
    """
    Manages Snowflake stored procedure interactions for the Data Exchange Agent.

    This class handles communication with the Snowflake stored procedure, including:
    - Loading Snowflake stored procedure configuration
    - Retrieving tasks from the Snowflake stored procedure
    - Updating task status and details

    The Snowflake stored procedure configuration is loaded from a TOML configuration file
    and used to configure the Snowflake stored procedure.
    """

    AGENT_TYPE = "data-exchange-agent"
    MAX_TASKS_PER_FETCH = 1
    PULL_TASKS_SP_NAME = "PULL_TASKS"
    COMPLETE_TASK_SP_NAME = "COMPLETE_TASK"
    FAIL_TASK_SP_NAME = "FAIL_TASK"
    COMPLETED_TASK_STATUS_KEY = COMPLETE_TASK_SP_NAME  # The SP returns a scalar value, so key's result is the SP name
    FAILED_TASK_STATUS_KEY = FAIL_TASK_SP_NAME  # The SP returns a scalar value, so key's result is the SP name

    @inject
    def __init__(
        self,
        program_config: ConfigManager = Provide[container_keys.PROGRAM_CONFIG],
        snowflake_datasource: SnowflakeDataSource = Provide[
            container_keys.SF_DATA_SOURCE
        ],
    ) -> None:
        """
        Initialize the SnowflakeStoredProcedureTaskSourceAdapter.

        Sets up the configuration attribute and loads the configuration from configuration.
        Registers a cleanup handler to close connections gracefully on program exit.
        """
        super().__init__()
        self._program_config = program_config

        try:
            self.agent_id = program_config[f"{config_keys.APPLICATION__AGENT_ID}"]
        except (KeyError, AttributeError) as e:
            raise custom_exceptions.ConfigurationError(
                "Agent ID is missing or incomplete."
            ) from e

        # Use the shared, thread-safe SnowflakeDataSource singleton.
        # Each worker thread will get its own connection via thread-local storage.
        self.snowflake_datasource = snowflake_datasource
        self.affinity: str | None = program_config.get(
            config_keys.APPLICATION__AFFINITY, None
        )
        self._qualified_data_migration_metadata_schema = (
            qualified_data_migration_metadata_schema(program_config)
        )

    def get_tasks(self) -> list[dict]:
        """
        Retrieve tasks from the Task Source.

        Makes a GET request to fetch tasks from the Task Source.

        Returns:
            list[dict]: List of task dictionaries from the Task Source response

        Raises:
            Exception: If the Task Source request fails

        """
        qm = self._qualified_data_migration_metadata_schema
        pull_tasks_statement = f"CALL {qm}.{self.PULL_TASKS_SP_NAME}(%s, %s, %s, %s)"
        params = (
            self.agent_id,
            self.AGENT_TYPE,
            self.affinity,
            self.MAX_TASKS_PER_FETCH,
        )
        try:
            # Reuse a persistent connection to avoid reconnecting on every poll.
            # The connection is only (re)created if it's closed or not yet established.
            if self.snowflake_datasource.is_closed():
                self.snowflake_datasource.create_connection()

            raw_results = self.snowflake_datasource.execute_statement(
                pull_tasks_statement, params
            )

            tasks_list = []
            for raw_task in raw_results:
                raw_task_payload = json.loads(raw_task["PAYLOAD"])
                task_kind = raw_task_payload.get(
                    task_keys.KIND, task_keys.DATA_MOVEMENT
                )

                if task_kind == task_keys.DATA_VALIDATION:
                    task_dict = {
                        task_keys.ID: raw_task["ID"],
                        task_keys.KIND: task_keys.DATA_VALIDATION,
                        task_keys.PAYLOAD: raw_task_payload,
                    }
                    tasks_list.append(task_dict)
                else:
                    task_obj = Task(
                        id=raw_task["ID"],
                        engine=raw_task_payload["source_type"],
                        database=raw_task_payload["database"],
                        schema=raw_task_payload["schema"],
                        statement=raw_task_payload["statement_location_id"],
                        target_type=raw_task_payload["target_type"],
                        target_id=raw_task_payload["target_id"],
                        suggested_batch_size=raw_task_payload.get(
                            "suggested_batch_size"
                        ),
                        is_bulk_operation=raw_task_payload.get(
                            "is_bulk_operation", False
                        ),
                        use_database_command=raw_task_payload.get(
                            "use_database_command"
                        ),
                    )
                    tasks_list.append(task_obj.to_dict())

            return tasks_list
        except ProgrammingError as e:
            if e.errno == TRANSIENT_SF_TRANSACTION_ERROR_CODE:
                logger.debug(
                    "Transient transaction conflict in PULL_TASKS (concurrent workers contention). Will retry on next cycle."
                )
                return []
            if e.errno == SF_UNKNOWN_FUNCTION_ERROR_CODE:
                logger.warning(
                    "Stored procedure not found (schema may still be deploying). Will retry on next cycle."
                )
                return []
            self.snowflake_datasource.close_connection()
            raise Exception(
                f"Snowflake error {e.errno} ({e.sqlstate}) while pulling tasks: {pull_tasks_statement}."
            ) from e
        except Exception as e:
            self.snowflake_datasource.close_connection()
            raise Exception(
                f"Failed to pull tasks executing statement: {pull_tasks_statement}."
            ) from e

    def complete_task(self, task_id: str) -> None:
        """
        Mark a task as completed.

        Args:
            task_id: The identifier of the task to mark as completed

        """
        qm = self._qualified_data_migration_metadata_schema
        complete_task_statement = f"CALL {qm}.{self.COMPLETE_TASK_SP_NAME}(%s, %s)"
        params = (self.agent_id, int(task_id))
        raw_result: dict | None = None
        try:
            if self.snowflake_datasource.is_closed():
                self.snowflake_datasource.create_connection()

            raw_result = next(
                iter(
                    self.snowflake_datasource.execute_statement(
                        complete_task_statement, params
                    )
                ),
                None,
            )
        except Exception as e:
            self.snowflake_datasource.close_connection()
            detail = f"Failed to mark task '{task_id}' as completed executing statement: {complete_task_statement}."
            raise Exception(detail) from e

        if raw_result is None or not raw_result.get(
            self.COMPLETED_TASK_STATUS_KEY, None
        ):
            detail = f"Failed to mark task '{task_id}' as completed executing statement: {complete_task_statement}."
            raise Exception(detail)

    def fail_task(self, task_id: str, error_message: str | None = None) -> None:
        """
        Mark a task as failed.

        Args:
            task_id: The identifier of the task to mark as failed
            error_message: Optional error message describing the failure

        """
        actual_error_message = (
            error_message if error_message else "No error message provided."
        )
        qm = self._qualified_data_migration_metadata_schema
        fail_task_statement = f"CALL {qm}.{self.FAIL_TASK_SP_NAME}(%s, %s, %s)"
        params = (int(task_id), self.agent_id, actual_error_message)

        try:
            if self.snowflake_datasource.is_closed():
                self.snowflake_datasource.create_connection()

            raw_result = next(
                iter(
                    self.snowflake_datasource.execute_statement(
                        fail_task_statement, params
                    )
                ),
                None,
            )
        except Exception as e:
            self.snowflake_datasource.close_connection()
            raise Exception(
                f"Failed to mark task '{task_id}' as failed executing statement: {fail_task_statement}."
            ) from e

        if raw_result is None or not raw_result.get(self.FAILED_TASK_STATUS_KEY, None):
            raise Exception(
                f"Failed to mark task '{task_id}' as failed executing statement: {fail_task_statement}."
            )
