"""
Lease refresher for maintaining task leases.

This module provides the LeaseRefresher class which runs a background thread
to periodically refresh task leases by calling the REFRESH_LEASES stored procedure.
"""

import threading

from dependency_injector.wiring import Provide, inject

from data_exchange_agent import custom_exceptions
from data_exchange_agent.config.manager import ConfigManager
from data_exchange_agent.constants import config_keys, container_keys
from data_exchange_agent.constants.config_defaults import (
    DEFAULT__APPLICATION__LEASE_REFRESH_INTERVAL,
)
from data_exchange_agent.data_sources.sf_connection import SnowflakeDataSource
from data_exchange_agent.snowflake_metadata_locations import (
    qualified_data_migration_metadata_schema,
)
from data_exchange_agent.utils.sf_logger import SFLogger


class LeaseRefresher:
    """
    Background thread that periodically refreshes task leases.

    This class ensures that tasks being processed by workers don't expire
    by periodically calling the REFRESH_LEASES stored procedure in Snowflake.

    Attributes:
        agent_id (str): Agent ID used to identify which leases to refresh
        refresh_interval (int): Interval in seconds between lease refresh calls
        _stop_event (threading.Event): Event to signal the thread to stop
        _thread (threading.Thread | None): The background refresh thread

    """

    REFRESH_LEASES_SP_NAME = "REFRESH_LEASES"

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

        Args:
            logger: Logger instance for logging messages
            program_config: Program configuration containing connection and agent settings
            snowflake_datasource: Shared Snowflake datasource (thread-safe, uses thread-local connections)

        Raises:
            ConfigurationError: If required configuration is missing

        """
        self.logger = logger

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

        self.refresh_interval = program_config.get(
            config_keys.APPLICATION__LEASE_REFRESH_INTERVAL,
            DEFAULT__APPLICATION__LEASE_REFRESH_INTERVAL,
        )
        self._qualified_data_migration_metadata_schema = (
            qualified_data_migration_metadata_schema(program_config)
        )

        # Use the shared, thread-safe SnowflakeDataSource singleton.
        # The LeaseRefresher thread will get its own connection via thread-local storage.
        self.snowflake_datasource = snowflake_datasource
        self._stop_event = threading.Event()
        self._thread: threading.Thread | None = None

    def start(self) -> None:
        """
        Start the lease refresh background thread.

        Creates and starts a daemon thread that will periodically refresh leases.
        If the thread is already running, this method does nothing.
        """
        if self._thread is not None and self._thread.is_alive():
            self.logger.warning("Lease refresher is already running.")
            return

        self._stop_event.clear()
        self._thread = threading.Thread(
            target=self._refresh_loop,
            daemon=True,
            name="LeaseRefresher",
        )
        self._thread.start()
        self.logger.info(f"Agent starting with UUID (consumer_id): {self.agent_id}")
        self.logger.info(
            f"Lease refresher started (interval: {self.refresh_interval}s)"
        )

    def stop(self) -> None:
        """
        Stop the lease refresh background thread.

        Signals the background thread to stop and waits for it to finish.
        """
        if self._thread is None or not self._thread.is_alive():
            return

        self.logger.info("Stopping lease refresher...")
        self._stop_event.set()
        self._thread.join(timeout=5)

        if self._thread.is_alive():
            self.logger.warning("Lease refresher thread did not stop gracefully.")
        else:
            self.logger.info("Lease refresher stopped.")

        self._thread = None

        # Close the Snowflake connection used by this thread
        self.snowflake_datasource.close_connection()

    def _refresh_loop(self) -> None:
        """
        Execute a loop that periodically refreshes leases.

        Runs until the stop event is set, calling the refresh_leases stored procedure
        at the configured interval.
        """
        while not self._stop_event.is_set():
            try:
                self._refresh_leases()
            except Exception as e:
                self.logger.error("Failed to refresh leases.", exception=e)

            # Wait for the interval or until stop is signaled
            self._stop_event.wait(timeout=self.refresh_interval)

    def _refresh_leases(self) -> None:
        """
        Call the REFRESH_LEASES stored procedure.

        This updates the leasing_time for all tasks currently leased to this agent,
        preventing them from being reclaimed by other agents.

        Uses a persistent connection to avoid the overhead of creating a new connection
        on each refresh cycle. The connection is only recreated if it's closed or fails.
        """
        qm = self._qualified_data_migration_metadata_schema
        refresh_statement = f"CALL {qm}.{self.REFRESH_LEASES_SP_NAME}(%s)"
        params = (self.agent_id,)

        try:
            # Ensure connection exists (creates one if closed or not yet established)
            if self.snowflake_datasource.is_closed():
                self.snowflake_datasource.create_connection()

            result = next(
                iter(
                    self.snowflake_datasource.execute_statement(
                        refresh_statement, params
                    )
                ),
                None,
            )
            if result:
                # The stored procedure returns a boolean indicating if any rows were updated
                leases_refreshed = result.get(self.REFRESH_LEASES_SP_NAME, False)
                if leases_refreshed:
                    self.logger.debug(f"Leases refreshed for agent '{self.agent_id}'")
                else:
                    self.logger.debug(
                        f"No active leases found for agent '{self.agent_id}'"
                    )
        except Exception as e:
            # Close the potentially broken connection so it gets recreated on the next attempt
            self.snowflake_datasource.close_connection()
            raise Exception(
                f"Failed to refresh leases executing statement: {refresh_statement}."
            ) from e
