import logging
import duckdb
import pandas as pd
import polars as pl
import pyarrow as pa
import re
from coded_flows.types import Union, List, DataFrame, ArrowTable, Str, Bool
from coded_flows.utils import CodedFlowsLogger

logger = CodedFlowsLogger(name="Delete Columns", level=logging.INFO)


def _coalesce(*values):
    return next((v for v in values if v is not None), None)


def _sanitize_identifier(identifier):
    """
    Sanitize SQL identifier by escaping special characters.
    Handles double quotes and other problematic characters.
    """
    return identifier.replace('"', '""')


def delete_columns(
    data: Union[DataFrame, ArrowTable],
    columns: List = None,
    regex_pattern: Str = None,
    options=None,
) -> Union[DataFrame, ArrowTable]:
    options = options or {}
    verbose = options.get("verbose", True)
    columns = _coalesce(columns, options.get("columns", []))
    regex_pattern = _coalesce(regex_pattern, options.get("regex_pattern", ""))
    output_format = options.get("output_format", "pandas")
    safe_mode = options.get("safe_mode", False)
    result = None
    try:
        selection_mode = None
        if regex_pattern:
            selection_mode = "regex"
        elif columns and len(columns) > 0:
            selection_mode = "list"
        if selection_mode == "list" and (not isinstance(columns, list)):
            verbose and logger.error(f"Invalid columns format! Expected a list.")
            raise ValueError("Columns must be provided as a list!")
        verbose and logger.info(
            f"Starting column delete operation. Mode: {selection_mode}."
        )
        data_type = None
        if isinstance(data, pd.DataFrame):
            data_type = "pandas"
        elif isinstance(data, pl.DataFrame):
            data_type = "polars"
        elif isinstance(data, (pa.Table, pa.lib.Table)):
            data_type = "arrow"
        if data_type is None:
            verbose and logger.error(
                f"Input data must be a pandas DataFrame, Polars DataFrame, or Arrow Table"
            )
            raise ValueError(
                "Input data must be a pandas DataFrame, Polars DataFrame, or Arrow Table"
            )
        verbose and logger.info(f"Detected input format: {data_type}.")
        conn = duckdb.connect(":memory:")
        conn.register("input_table", data)
        column_info = conn.execute("DESCRIBE input_table").fetchall()
        all_columns = [col[0] for col in column_info]
        columns_to_delete = set()
        skipped_count = 0
        if selection_mode == "regex":
            try:
                pattern = re.compile(regex_pattern)
                matched_cols = [col for col in all_columns if pattern.search(col)]
                columns_to_delete.update(matched_cols)
                verbose and logger.info(
                    f"Regex pattern '{regex_pattern}' matched {len(matched_cols)} columns to delete."
                )
            except re.error:
                conn.close()
                verbose and logger.error(f"Invalid regex pattern provided.")
                raise ValueError(f"Invalid regex pattern: {regex_pattern}")
        elif selection_mode == "list":
            if not safe_mode:
                missing_columns = [col for col in columns if col not in all_columns]
                if missing_columns:
                    verbose and logger.error(
                        f"Columns not found in data: {missing_columns}"
                    )
                    conn.close()
                    raise ValueError(f"Columns not found in data: {missing_columns}")
            for col in columns:
                if col in all_columns:
                    columns_to_delete.add(col)
                elif safe_mode:
                    verbose and logger.warning(
                        f"Safe mode: Skipping non-existent column: {col}"
                    )
                    skipped_count += 1
        select_parts = []
        kept_count = 0
        deleted_count = 0
        for col in all_columns:
            if col not in columns_to_delete:
                sanitized_col = _sanitize_identifier(col)
                select_parts.append(f'"{sanitized_col}"')
                kept_count += 1
            else:
                deleted_count += 1
        if not select_parts:
            verbose and logger.error(
                f"All columns would be deleted! Result would be empty."
            )
            conn.close()
            raise ValueError(
                "Cannot delete all columns! At least one column must remain."
            )
        select_clause = ", ".join(select_parts)
        query = f"SELECT {select_clause} FROM input_table"
        verbose and logger.info(f"Executing query to delete columns.")
        if output_format == "pandas":
            result = conn.execute(query).df()
            verbose and logger.info(f"Converted result to pandas DataFrame.")
        elif output_format == "polars":
            result = conn.execute(query).pl()
            verbose and logger.info(f"Converted result to Polars DataFrame.")
        elif output_format == "arrow":
            result = conn.execute(query).fetch_arrow_table()
            verbose and logger.info(f"Converted result to Arrow Table.")
        else:
            conn.close()
            verbose and logger.error(f"Unsupported output format: {output_format}")
            raise ValueError(f"Unsupported output format: {output_format}")
        conn.close()
        verbose and logger.info(
            f"Operation completed. Deleted {deleted_count} columns, kept {kept_count} columns."
        )
    except Exception as e:
        verbose and logger.error(f"Error during column delete operation.")
        raise
    return result