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

logger = CodedFlowsLogger(name="Split Target Column", level=logging.INFO)


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


def _sanitize_identifier(identifier):
    """
    Sanitize SQL identifier by escaping special characters.
    """
    return identifier.replace('"', '""')


def split_target_column(
    data: Union[DataFrame, ArrowTable], target_column: Str = None, options=None
) -> Tuple[DataFrame, DataSeries]:
    brick_display_name = "Split Target Column"
    options = options or {}
    verbose = options.get("verbose", True)
    target_col_name = _coalesce(target_column, options.get("target_column"))
    output_format = options.get("output_format", "pandas")
    X = None
    y = None
    if not target_col_name:
        verbose and logger.error(f"[{brick_display_name}] No target column specified.")
        raise ValueError("A target column name must be provided.")
    try:
        verbose and logger.info(
            f"[{brick_display_name}] Starting split operation. Target: '{target_col_name}'"
        )
        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:
            raise ValueError(
                "Input data must be a pandas DataFrame, Polars DataFrame, or Arrow Table"
            )
        verbose and logger.info(
            f"[{brick_display_name}] 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]
        if target_col_name not in all_columns:
            conn.close()
            raise ValueError(f"Target column '{target_col_name}' not found in dataset.")
        sanitized_target = _sanitize_identifier(target_col_name)
        query_X = f'SELECT * EXCLUDE ("{sanitized_target}") FROM input_table'
        query_y = f'SELECT "{sanitized_target}" FROM input_table'
        verbose and logger.info(f"[{brick_display_name}] Executing split queries.")
        if output_format == "pandas":
            X = conn.execute(query_X).df()
            y_df = conn.execute(query_y).df()
            y = y_df.iloc[:, 0]
            verbose and logger.info(
                f"[{brick_display_name}] Converted result to pandas DataFrame (X) and Series (y)."
            )
        elif output_format == "polars":
            X = conn.execute(query_X).pl()
            y_df = conn.execute(query_y).pl()
            y = y_df.to_series()
            verbose and logger.info(
                f"[{brick_display_name}] Converted result to Polars DataFrame (X) and Series (y)."
            )
        else:
            conn.close()
            raise ValueError(
                f"Unsupported output format: {output_format}. Must be 'pandas' or 'polars'."
            )
        conn.close()
        verbose and logger.info(
            f"[{brick_display_name}] Split complete. X shape: {X.shape}, y shape: {y.shape}"
        )
    except Exception as e:
        verbose and logger.error(f"[{brick_display_name}] Error during split operation")
        raise
    return (X, y)