"""DuckDB utility module."""

from __future__ import annotations

from pathlib import Path

from duckdb import (
    DuckDBPyConnection as DuckDBPyCursor,
)

from src.connectors.snowflake import SnowflakeConfig
from src.enums import FileType
from src.errors import FileParsingError, InvalidFileTypeError
from src.schemas import FileMetadata
from src.utils.db_utils import get_select_clause


def add_s3_secret(
    cursor: DuckDBPyCursor,
    region: str = 'us-east-1',
    secret_name: str = 'my_s3_secret',
    replace: bool = False,
) -> str:
    """Add S3 credentials/configuration as a DuckDB secret.

    Args:
        cursor: Active DuckDB cursor.
        region: AWS region for the S3 bucket.
        secret_name: Optional secret name. Defaults to 'my_s3_secret'.
        replace: Whether to use 'CREATE OR REPLACE'. Defaults to False.
    """
    or_replace = 'OR REPLACE' if replace else ''
    cursor.execute(
        f"""
        CREATE {or_replace} SECRET {secret_name} (
            TYPE S3,
            PROVIDER credential_chain,
            REGION '{region}'
        )
        """
    )
    return secret_name


def add_snowflake_secret(
    cursor: DuckDBPyCursor,
    config: SnowflakeConfig,
    secret_name: str = 'my_snowflake_secret',
    replace: bool = False,
) -> str:
    """Add Snowflake credentials as a DuckDB secret for secure connection.

    Args:
        cursor: Active DuckDB cursor.
        config: Snowflake configuration with credentials.
        secret_name: Optional secret name. Defaults to 'my_snowflake_secret'.
        replace: Whether to use 'CREATE OR REPLACE'. Defaults to False.
    """
    params = {
        'ACCOUNT': config.account,
        'USER': config.user,
        'ROLE': config.role,
        'DATABASE': config.database,
        'WAREHOUSE': config.warehouse,
    }

    if config.private_key:
        params['PRIVATE_KEY'] = config.private_key
    elif config.private_key_path:
        params['PRIVATE_KEY_PATH'] = config.private_key_path

    if config.private_key_passphrase:
        params['PRIVATE_KEY_PASSPHRASE'] = config.private_key_passphrase

    cmd = 'CREATE OR REPLACE' if replace else 'CREATE'
    key_vals = [f"{key} '{value}'" for key, value in params.items()]
    cursor.execute(
        f"""
        {cmd} SECRET {secret_name} (
            TYPE snowflake,
            AUTH_TYPE key_pair,
            {', '.join(key_vals)}
        )
        """
    )

    return secret_name


# -- File Inspection --


def get_file_columns(cursor: DuckDBPyCursor, meta: FileMetadata) -> set[str]:
    """Get column names from a file.

    Args:
        cursor: Active DuckDB connection cursor.
        meta: Metadata about the file including path, type, and encoding.

    Returns:
        set[str]: A set of column names.

    Raises:
        InvalidFileTypeError: If file type is not supported.
        FileParsingError: If there is an error loading the file.
    """
    read_query = _get_read_query(meta)

    try:
        result = cursor.execute(
            f'DESCRIBE SELECT * FROM {read_query}',
            [str(meta.file_path)],
        ).fetchall()

        return {column_meta[0] for column_meta in result}
    except Exception as e:
        raise FileParsingError(f'Failed to parse {meta.file_type} columns: {e}') from e


# --- File -> Table ---


def create_table_from_file(
    cursor: DuckDBPyCursor,
    table_name: str,
    meta: FileMetadata,
    alias_map: dict[str, str | None] | None = None,
) -> None:
    """Create a DuckDB table from a file (CSV, Parquet, Excel).

    Args:
        cursor: Active DuckDB connection cursor.
        table_name: The name to assign to the DuckDB table.
        meta: Metadata about the file.
        alias_map: Optional mapping from source columns to target columns.
            If provided, only mapped columns are loaded and renamed.

    Raises:
        InvalidFileTypeError: If the file type is not supported or gzipped XLSX.
        FileParsingError: If there is an error loading the file into DuckDB.
    """
    read_query = _get_read_query(meta)

    try:
        select_clause = get_select_clause(alias_map)
        cursor.execute(
            f"""
                CREATE OR REPLACE TABLE "{table_name}" AS
                SELECT {select_clause}
                FROM {read_query}
            """,
            [str(meta.file_path)],
        )
    except Exception as e:
        raise FileParsingError(f'Failed to load {meta.file_type} file: {e}') from e


def _get_read_query(meta: FileMetadata) -> str:
    """Build the read_csv function call string.

    Args:
        meta: Metadata about the file.

    Returns:
        str: The read function call string.
    """
    if meta.file_type == FileType.PQT:
        return 'read_parquet(?)'

    if meta.file_type == FileType.XLSX:
        if meta.gzipped:
            raise InvalidFileTypeError('Gzipped XLSX files are not supported.')
        return 'read_xlsx(?, ALL_VARCHAR = TRUE, HEADER = TRUE)'

    if meta.file_type == FileType.CSV:
        args = [
            '?',
            'ALL_VARCHAR = TRUE',
            'HEADER = TRUE',
            'NULL_PADDING = TRUE',
            'SAMPLE_SIZE = 64',
        ]

        if meta.gzipped:
            args.append("COMPRESSION = 'gzip'")

        enc = _sanitize_encoding(meta.encoding)
        if enc:
            args.append(f"ENCODING = '{enc}'")

        return f'read_csv({", ".join(args)})'

    raise InvalidFileTypeError(f'Unsupported file type: {meta.file_type}')


def _sanitize_encoding(encoding: str | None) -> str | None:
    """Sanitize encoding string for DuckDB compatibility.

    DuckDB versions 1.4.2 and 1.4.3 have a bug where explicitly
    specifying UTF-8 encoding causes the operation to hang. This
    method filters out UTF-8 variants and normalizes other encodings.

    Args:
        encoding: The encoding string to sanitize. Can be None.

    Returns:
        The lowercased encoding string if it's not a UTF-8 variant,
        or None if the input is None, empty, or a UTF-8 variant
        (utf-8, utf-8-sig, utf8).
    """
    if not encoding:
        return None

    encoding = encoding.strip().replace('_', '-').lower()
    return None if encoding in ['utf-8', 'utf-8-sig', 'utf8'] else encoding


# --- Table Export ---


def create_csv_from_table(
    cursor: DuckDBPyCursor,
    table_name: str,
    file_path: str | Path,
    columns: list[str] | None = None,
    rowid: bool = False,
    gzipped: bool = False,
) -> None:
    """Export a DuckDB table to a CSV file.

    Exports table data to a CSV file with optional gzip compression, allowing
    customization of output columns including prefix columns, row IDs, and all
    table columns.

    Args:
        cursor: Active DuckDB connection cursor.
        table_name: The name of the table to export.
        file_path: The path to the output CSV file.
        columns: The column names to export, or all if not given.
        rowid: Whether to include a row_id column (1-indexed) in the output.
            The row_id will be inserted before the table columns. Default: False.
        gzipped: Whether to gzip-compress the output file. Default: False.
            Note: MySQL LOAD DATA (both LOCAL INFILE and FROM S3) does not
            support gzippeded CSV files.
    """
    select_query = _build_export_select_query(columns, rowid)
    compression = 'COMPRESSION GZIP,' if gzipped else ''

    cursor.execute(
        f"""
        COPY (
            SELECT {select_query}
            FROM "{table_name}"
        ) TO "{file_path}" (
            {compression}
            FORMAT CSV,
            HEADER TRUE
        )
        """
    )


def create_parquet_from_query(
    cursor: DuckDBPyCursor,
    query: str,
    file_path: str | Path,
    gzipped: bool = False,
) -> None:
    """Export a DuckDB table to a parquet file.

    Exports table data to a parquet file with optional gzip compression, allowing
    customization of output columns including prefix columns, row IDs, and all
    table columns.

    Args:
        cursor: Active DuckDB connection cursor.
        query: The query to export.
        file_path: The path to the output parquet file.
        gzipped: Whether to gzip-compress the output file. Default: False.
            Note: MySQL LOAD DATA (both LOCAL INFILE and FROM S3) does not
            support gzippeded parquet files.
    """
    compression = 'COMPRESSION GZIP,' if gzipped else ''

    cursor.execute(
        f"""
        COPY (
            SELECT *
            FROM ({query})
        ) TO "{file_path}" (
            {compression}
            FORMAT PARQUET
        )
        """
    )


def upload_csv_from_table(
    cursor: DuckDBPyCursor,
    table_name: str,
    s3_bucket: str,
    s3_key: str,
    columns: list[str] | None = None,
    rowid: bool = False,
) -> None:
    """Export a DuckDB table to a CSV file in S3.

    Exports table data to a CSV file with optional gzip compression, allowing
    customization of output columns including prefix columns, row IDs, and all
    table columns.

    Args:
        cursor: The DuckDB cursor.
        table_name: The name of the table to export.
        s3_bucket: The S3 bucket.
        s3_key: The S3 key.
        columns: The column names to export, or all if not given.
        rowid: Whether to include a row_id column (1-indexed) in the output.
            The row_id will be inserted before the table columns. Default: False.
    """
    select_query = _build_export_select_query(columns, rowid)
    cursor.execute(
        f"""
        COPY (
            SELECT {select_query}
            FROM "{table_name}"
        ) TO 's3://{s3_bucket}/{s3_key}' (
            FORMAT CSV,
            HEADER TRUE
        )
        """
    )


def _build_export_select_query(columns: list[str] | None, rowid: bool) -> str:
    """Build SELECT clause for table export.

    Args:
        columns: Optional list of column names to export.
        rowid: Whether to include a row_id column.

    Returns:
        str: The SELECT clause.
    """
    parts: list[str] = []

    if rowid:
        parts.append('rowid + 1 AS row_id')

    if columns:
        parts.extend([f'"{col}"' for col in columns])
    else:
        parts.append('*')

    return ', '.join(parts)


# --- Table Inspection ---


def get_table_string_columns(cursor: DuckDBPyCursor, table_name: str) -> list[str]:
    """Get the column names of a DuckDB table that are text-based.

    Args:
        cursor: Active DuckDB connection cursor.
        table_name: The name of the table.

    Returns:
        list[str]: A list of text-based column names, ordered by position.
    """
    column_name_result = cursor.execute(
        """
        SELECT column_name
        FROM information_schema.columns
        WHERE table_name = ?
          AND data_type IN ('VARCHAR', 'BPCHAR', 'TEXT', 'STRING')
        ORDER BY ordinal_position ASC
        """,
        [table_name],
    ).fetchall()
    return [row[0] for row in column_name_result]


# --- Table Modification ---


def normalize_table_rows(
    cursor: DuckDBPyCursor,
    table_name: str,
    max_len: int | None = None,
) -> None:
    r"""Normalize text columns by trimming and removing control characters.

    1. Identifies string-type columns (ignores INT, DOUBLE, DATE).
    2. Removes control characters that break MySQL/UTF-8 (0x00-0x1F, 0x7F).
    3. Preserves legitimate whitespace (TAB, LF, CR).
    4. Trims leading/trailing whitespace.
    5. Preserves NULL values (does not convert them to empty strings).

    Removes:
    1. \x00 - \x08 : NULL to Backspace
    2. \x0B - \x0C : Vertical Tab and Form Feed
    3. \x0E - \x1F : Shift Out to Unit Separator
    4. \x7F         : Delete

    Preserves:
    1. \x09 : TAB
    2. \x0A : LF
    3. \x0D : CR

    Args:
        cursor: Active DuckDB connection cursor.
        table_name: The name of the table.
        max_len: Optional max length to truncate string columns to.
    """
    if max_len is not None and max_len < 0:
        raise ValueError(f'max_len must be nonnegative. Received: {max_len}')

    string_columns = get_table_string_columns(cursor, table_name)
    if not string_columns:
        return

    chars = r'[\x00-\x08\x0B-\x0C\x0E-\x1F\x7F]+'
    cmd_tpl = f"TRIM(REGEXP_REPLACE(\"{{col}}\", '{chars}', '', 'g'))"
    cmd_tpl = f'LEFT({cmd_tpl}, {max_len})' if max_len is not None else cmd_tpl

    parts = [f'"{col}" = {cmd_tpl.format(col=col)}' for col in string_columns]
    cursor.execute(f'UPDATE "{table_name}" SET {", ".join(parts)}')