"""DuckDB connection adapter.

Implements the DBConnection protocol for in-memory analytics and data validation.
"""

import os
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path

import duckdb
import psutil
from abacus_common_logic.constants.math import Bytes
from duckdb import DuckDBPyConnection, DuckDBPyConnection as DuckDBPyCursor
from pydantic import BaseModel, Field

from src.connectors.db import Connection
from src.infra.log import logger
from src.utils.file_utils import get_free_disk_space, get_num_cpu_workers


class DuckDBConfig(BaseModel):
    """Base model for DB configs."""

    file_path: str | Path | None = Field(
        default=None, description='Path to the database file. If None, uses in-memory.'
    )
    max_memory_mb: int | None = Field(
        default=None, description='Optional memory limit in MB.'
    )
    max_temp_directory_mb: int | None = Field(
        default=None, description='Optional max temp directory size in MB.'
    )
    max_threads: int | None = Field(
        default=None, description='Optional maximum number of threads.'
    )


class DuckDBConnection(Connection):
    """DuckDB database connection adapter."""

    def __init__(self, conn: DuckDBPyConnection):
        """Initialize the adapter with a raw duckdb connection.

        Args:
            conn: An active duckdb connection.
        """
        self._conn = conn

    def close(self) -> None:
        """Close the connection."""
        self._conn.close()

    def commit(self) -> None:
        """Commit the current transaction to the database."""
        self._conn.commit()

    def cursor(self) -> DuckDBPyCursor:
        """Return a raw duckdb cursor for low-level use."""
        return self._conn.cursor()

    # pd.read_sql(sql, self._conn, params=params, dtype_backend='pyarrow')

    def rollback(self) -> None:
        """Roll back the current transaction, discarding changes."""
        self._conn.rollback()


class DuckDBConnectionFactory:
    """Factory for creating DuckDB database connections with configuration management."""

    def __init__(self, config: DuckDBConfig) -> None:
        """Initialize the connection factory with DuckDB configuration.

        Args:
            config: DuckDBConfig instance containing database connection params
        """
        self._config = config

    @contextmanager
    def connection(self) -> Generator[DuckDBConnection, None, None]:
        """Create and yield a DuckDB database connection as a context manager.

        The connection is automatically closed when the context exits.

        Yields:
            DuckDBConnection: DuckDB database adapter.

        Raises:
            duckdb.Error: If the connection attempt fails.
            IsADirectoryError: If the provided file_path points to a directory.
        """
        with duckdb_connection(self._config) as conn:
            yield DuckDBConnection(conn)


@contextmanager
def duckdb_connection(
    config: DuckDBConfig,
) -> Generator[DuckDBPyConnection, None, None]:
    """Context manager for DuckDB connection with resource limits.

    Args:
        config: DuckDB configuration object.

    Yields:
        DuckDBPyConnection: An initialized connection to the specified database.

    Raises:
        duckdb.Error: If the connection attempt fails.
        IsADirectoryError: If the provided file_path points to a directory.
    """
    # Sanitizing file
    file_path = config.file_path
    if file_path:
        if os.path.isdir(file_path):
            raise IsADirectoryError(
                f"Provided file_path '{file_path}' is a directory. "
                'DuckDB requires a file path.'
            )

        if os.path.isfile(file_path) and os.path.getsize(file_path) == 0:
            logger.info(f'Removing 0-byte placeholder file: {file_path}')
            os.remove(file_path)

        db_dir = os.path.dirname(file_path)
        os.makedirs(db_dir, exist_ok=True)

    conn: DuckDBPyConnection | None = None
    try:
        database = str(file_path) if file_path else ':memory:'
        logger.info(f'Connecting to DuckDB: {database}')

        # Lazy import to avoid circular dependency
        from config import config as _config

        # Baked home directory with installed extensions (read-only)
        logger.info(f'- home_directory: {_config.duckdb.HOME_DIR}')
        os.makedirs(_config.duckdb.HOME_DIR, exist_ok=True)

        # Ephemeral swap directory for spilling data from RAM if needed
        logger.info(f'- temp_directory: {_config.duckdb.TEMP_DIR}')
        os.makedirs(_config.duckdb.TEMP_DIR, exist_ok=True)

        # Max RAM usage
        max_memory_mb = config.max_memory_mb
        if not max_memory_mb:
            total_ram = psutil.virtual_memory().total / Bytes.MB
            max_memory_mb = int(_config.duckdb.MEM_PCT * total_ram)

        # Max memory for temp files
        max_temp_directory_mb = config.max_temp_directory_mb
        if not max_temp_directory_mb:
            lambda_temp_dir_mb = get_free_disk_space(_config.duckdb.TEMP_DIR) / Bytes.MB
            max_temp_directory_mb = int(_config.duckdb.TEMP_PCT * lambda_temp_dir_mb)

        # Max threads
        max_threads = config.max_threads
        if not max_threads:
            max_threads = get_num_cpu_workers()
        if _config.duckdb.MAX_THREADS:
            max_threads = min(_config.duckdb.MAX_THREADS, max_threads)
        max_threads = max(1, max_threads)

        conn = duckdb.connect(
            database=database,
            config={
                'autoinstall_known_extensions': False,
                'autoload_known_extensions': False,
                'max_memory': f'{max_memory_mb}MB',
                'max_temp_directory_size': f'{max_temp_directory_mb}MB',
                'preserve_insertion_order': False,
                'threads': max_threads,
            },
        )

        conn.execute(f"SET home_directory = '{_config.duckdb.HOME_DIR}'")
        conn.execute(f"SET temp_directory = '{_config.duckdb.TEMP_DIR}'")
        conn.execute("SET logging_level = 'WARN'")

        extensions = {'aws', 'encodings', 'excel', 'httpfs', 'snowflake'}
        for ext in extensions:
            # Skips install. Assumes extensions are baked in.
            conn.execute(f'LOAD {ext}')

        yield conn
    finally:
        if conn:
            conn.close()