"""
SQL operations module wrapping fullmetalalchemy and transmutation.

This module provides high-level functions for executing SQL operations
with automatic transaction management and rollback on errors.
"""

from __future__ import annotations

import contextlib
from typing import Any

import fullmetalalchemy as fa
import fullmetalalchemy.connection as fac
import pandas as pd
import transmutation as tm
from sqlalchemy import Engine

from pandalchemy.exceptions import TransactionError
from pandalchemy.execution_plan import ExecutionPlan, OperationType, SchemaChange
from pandalchemy.utils import (
    convert_numpy_types,
    convert_records_list,
    extract_primary_key_column,
    normalize_schema,
    pandas_dtype_to_python_type,
)


def pull_table(
    engine: Engine,
    table_name: str,
    schema: str | None = None,
    primary_key: str | list[str] | None = None,
    set_index: bool = True,
) -> pd.DataFrame:
    """
    Pull a table from the database into a DataFrame.

    Args:
        engine: SQLAlchemy engine
        table_name: Name of the table to pull
        schema: Optional schema name
        primary_key: Primary key column(s) to set as index
        set_index: Whether to set primary key as index (default True)

    Returns:
        DataFrame with PK as index if set_index=True and primary_key provided
    """
    df = fa.select.select_table_as_dataframe(
        table_name, connection=engine, schema=schema, primary_key=primary_key, set_index=set_index
    )

    # Fix type conversions after pulling from database
    # SQLite stores booleans as integers (0/1), and sometimes they come back as strings
    # Convert boolean-like columns back to proper booleans
    df = _fix_dataframe_types(df)

    return df


def _fix_dataframe_types(df: pd.DataFrame) -> pd.DataFrame:
    """
    Fix data types after pulling from database.

    SQLite stores booleans as integers, and some databases return strings.
    This function converts them back to proper types.
    Also fixes index types when they're numeric strings.

    Args:
        df: DataFrame with potentially incorrect types

    Returns:
        DataFrame with corrected types
    """
    if len(df) == 0:
        return df

    df = df.copy()

    # Fix index types first (for primary key lookup)
    if df.index.name and len(df.index) > 0:
        try:
            # Try to convert index to numeric if it looks numeric
            if df.index.dtype == "object":
                numeric_idx = pd.to_numeric(df.index, errors="coerce")
                if not numeric_idx.isna().all():
                    # All values are numeric - convert index
                    # Use Series to check if integers
                    numeric_series = pd.Series(numeric_idx)
                    if numeric_series.dropna().mod(1).eq(0).all():
                        # All integers
                        df.index = pd.Index(numeric_idx.astype("Int64"), name=df.index.name)
                    else:
                        # Floats
                        df.index = pd.Index(numeric_idx, name=df.index.name)
        except (ValueError, TypeError):
            pass

    # Fix MultiIndex types
    if isinstance(df.index, pd.MultiIndex) and len(df.index) > 0:
        new_levels = []
        for _i, level in enumerate(df.index.levels):
            try:
                if level.dtype == "object":
                    numeric_level = pd.to_numeric(level, errors="coerce")
                    if not numeric_level.isna().all():
                        # Use Series to check if integers
                        numeric_series = pd.Series(numeric_level)
                        if numeric_series.dropna().mod(1).eq(0).all():
                            new_levels.append(
                                pd.Index(numeric_level.astype("Int64"), name=level.name)
                            )
                        else:
                            new_levels.append(pd.Index(numeric_level, name=level.name))
                    else:
                        new_levels.append(level)
                else:
                    new_levels.append(level)
            except (ValueError, TypeError):
                new_levels.append(level)

        if len(new_levels) == len(df.index.levels):
            with contextlib.suppress(Exception):
                df.index = pd.MultiIndex.from_arrays(
                    [new_levels[i] for i in range(len(new_levels))], names=df.index.names
                )

    for col in df.columns:
        if len(df[col].dropna()) == 0:
            continue

        # Skip conversion if column is already boolean type - trust the data
        if df[col].dtype in ["bool", "boolean"]:
            continue

        # Check if column looks like boolean (contains only 0/1 or '0'/'1'/'True'/'False')
        unique_values = {str(v).lower() for v in df[col].dropna().unique()}
        boolean_like = unique_values.issubset({"0", "1", "true", "false", "0.0", "1.0", 0, 1})

        if boolean_like and len(unique_values) <= 3:  # Only boolean values
            # Convert to boolean
            try:
                if df[col].dtype == "object":
                    # String values - handle '0', '1', 'True', 'False', 'true', 'false'
                    def to_bool(val):
                        if pd.isna(val):
                            return pd.NA
                        s = str(val).lower()
                        if s in ("true", "1", "1.0"):
                            return True
                        elif s in ("false", "0", "0.0"):
                            return False
                        return val

                    df[col] = df[col].apply(to_bool).astype("boolean")
                elif df[col].dtype in ["int64", "Int64", "int32", "Int32", "float64", "Float64"]:
                    # Integer/float values (0/1) - be careful not to convert non-boolean columns
                    # Only convert if all values are exactly 0 or 1
                    unique_numeric = set(df[col].dropna().unique())
                    if unique_numeric.issubset({0, 1}):
                        df[col] = df[col].astype(bool)
                elif df[col].dtype == "bool":
                    # Already boolean - keep as is
                    pass
                elif df[col].dtype == "boolean":
                    # Already nullable boolean - keep as is
                    pass
            except (ValueError, TypeError):
                # If conversion fails, keep original
                pass
        # Convert numeric strings to integers/floats, but be conservative
        # Only convert if ALL non-null values are numeric and the column name suggests it's numeric
        elif df[col].dtype == "object" and col not in (
            df.index.names
            if hasattr(df.index, "names")
            else [df.index.name]
            if df.index.name
            else []
        ):
            try:
                # Check if column name suggests it should be a string (not numeric)
                # Columns with these words are typically strings: name, code, text, description, title, label, message, comment, note
                # Columns with these patterns are typically numeric: id, count, amount, price, quantity, value (context-dependent)
                string_indicators = {
                    "name",
                    "code",
                    "text",
                    "description",
                    "title",
                    "label",
                    "message",
                    "comment",
                    "note",
                    "address",
                    "email",
                    "phone",
                }
                col_lower = col.lower()

                # Special handling for columns that could be either string or numeric
                # "value" and "id" columns ending in "_id" are often numeric
                is_likely_numeric = (
                    col_lower.endswith("_id")
                    or col_lower == "value"
                    or "count" in col_lower
                    or "amount" in col_lower
                    or "price" in col_lower
                )

                # If column name suggests string and it's not likely numeric, skip conversion
                if (
                    any(indicator in col_lower for indicator in string_indicators)
                    and not is_likely_numeric
                ):
                    pass
                else:
                    # Try to convert numeric strings to integers
                    numeric_col = pd.to_numeric(df[col], errors="coerce")
                    if not numeric_col.isna().all():
                        # Check if all non-null values are integers (not floats)
                        non_null = numeric_col.dropna()
                        # If all non-null values are integers (whole numbers), convert to Int64
                        # Only convert if it looks intentional (most values are numeric)
                        if (
                            len(non_null) > 0
                            and non_null.mod(1).eq(0).all()
                            and len(non_null) / len(df[col]) > 0.8
                        ):  # At least 80% are numeric
                            df[col] = numeric_col.astype("Int64")
            except (ValueError, TypeError):
                pass

    return df


def get_primary_key(
    engine: Engine, table_name: str, schema: str | None = None
) -> str | list[str] | None:
    """
    Get the primary key column name(s) for a table.

    Args:
        engine: SQLAlchemy engine
        table_name: Name of the table
        schema: Optional schema name

    Returns:
        Primary key column name (str) for single-column PK,
        list of column names for composite PK,
        or None if no primary key exists
    """
    return fa.features.get_primary_key_names(table_name, engine, schema)


def execute_plan(
    engine: Engine,
    table_name: str,
    plan: ExecutionPlan,
    schema: str | None = None,
    primary_key: str | list[str] | None = None,
) -> None:
    """
    Execute a complete execution plan with transaction management.

    All operations are executed within a single transaction. If any operation
    fails, all changes are rolled back automatically.

    Args:
        engine: SQLAlchemy engine
        table_name: Name of the table to modify
        plan: The ExecutionPlan to execute
        schema: Optional schema name
        primary_key: Name of the primary key column(s) - single string or list for composite

    Raises:
        SQLAlchemyError: If any SQL operation fails
    """
    if not plan.has_changes():
        return

    # Separate schema changes from data changes
    # Schema changes must execute OUTSIDE any transaction to avoid metadata lock conflicts
    schema_changes = [
        step for step in plan.steps if step.operation_type == OperationType.SCHEMA_CHANGE
    ]
    data_changes = [
        step for step in plan.steps if step.operation_type != OperationType.SCHEMA_CHANGE
    ]

    # Execute schema changes first - outside data transaction to avoid metadata lock conflicts
    # Transmutation handles its own transactions when given engine parameter
    for step in schema_changes:
        _execute_schema_change(engine, table_name, step.data, schema)
        # Transmutation commits automatically when using engine parameter

    # Execute data changes within a transaction
    if data_changes:
        with engine.begin() as connection:
            try:
                for step in data_changes:
                    if step.operation_type == OperationType.DELETE:
                        _execute_deletes(connection, table_name, step.data, primary_key, schema)
                    elif step.operation_type == OperationType.UPDATE:
                        _execute_updates(connection, table_name, step.data, primary_key, schema)
                    elif step.operation_type == OperationType.INSERT:
                        _execute_inserts(connection, table_name, step.data, schema)
            except Exception as e:
                # Transaction will automatically rollback on exception
                raise TransactionError(
                    f"Failed to execute plan: {str(e)}",
                    details={"table": table_name, "schema": schema},
                ) from e


def _execute_schema_change(
    engine: Engine, table_name: str, schema_change: SchemaChange, schema: str | None = None
) -> None:
    """
    Execute a schema change operation using transmutation.

    Args:
        engine: SQLAlchemy engine
        table_name: Name of the table to modify
        schema_change: The schema change to apply
        schema: Optional schema name
    """
    # Transmutation supports both engine and connection parameters
    # Using engine ensures transmutation manages its own transactions correctly
    # Normalize schema: empty string -> None
    schema_normalized = normalize_schema(schema)

    # Determine MySQL-specific defaults
    is_mysql = engine.dialect.name == "mysql"
    default_varchar_length = 255 if is_mysql else None

    # Wrap transmutation calls to convert ValidationError to TransactionError
    try:
        if schema_change.change_type == "add_column":
            tm.add_column(
                table_name=table_name,
                column_name=schema_change.column_name,
                dtype=_pandas_dtype_to_python_type(schema_change.column_type),
                engine=engine,  # Engine parameter ensures transmutation manages transactions
                schema=schema_normalized,
                default_varchar_length=default_varchar_length,  # MySQL VARCHAR(255) default
                verify=False,  # Skip verification to avoid metadata lock issues
            )
        elif schema_change.change_type == "drop_column":
            tm.drop_column(
                table_name=table_name,
                col_name=schema_change.column_name,
                engine=engine,
                schema=schema_normalized,
                verify=False,
            )
        elif schema_change.change_type == "rename_column":
            if schema_change.new_column_name is None:
                raise ValueError("new_column_name is required for rename_column operation")
            # existing_type=None enables auto-detection for MySQL
            tm.rename_column(
                table_name=table_name,
                old_col_name=schema_change.column_name,
                new_col_name=schema_change.new_column_name,
                engine=engine,
                schema=schema_normalized,
                existing_type=None,  # Auto-detect existing type (especially for MySQL)
                verify=False,  # Skip verification to avoid metadata lock issues
            )
        elif schema_change.change_type == "alter_column_type":
            # Transmutation's alter_column can lose data when converting types (e.g., TEXT to INTEGER)
            # The execution plan should handle updating all rows with correct converted values
            # We just need to perform the schema change here
            tm.alter_column(
                table_name=table_name,
                column_name=schema_change.column_name,
                engine=engine,
                schema=schema_normalized,
                type_=_pandas_dtype_to_python_type(schema_change.new_column_type),
                verify=False,  # Skip verification to avoid metadata lock issues
            )
    except Exception as e:
        # Convert transmutation ValidationError and other errors to TransactionError
        if (
            "ValidationError" in type(e).__name__
            or "Column" in str(e)
            and ("does not exist" in str(e) or "not found" in str(e))
        ):
            raise TransactionError(
                f"Schema change failed: {str(e)}",
                details={
                    "table": table_name,
                    "schema": schema_normalized,
                    "change_type": schema_change.change_type,
                },
            ) from e
        # Re-raise other exceptions as-is
        raise


def _execute_deletes(
    connection: Any,
    table_name: str,
    delete_keys: list[Any],
    primary_key: str | list[str] | None,
    schema: str | None = None,
) -> None:
    """
    Execute delete operations.

    Args:
        connection: SQLAlchemy connection within a transaction
        table_name: Name of the table
        delete_keys: List of primary key values to delete (tuples for composite PKs)
        primary_key: Name of the primary key column(s)
        schema: Optional schema name
    """
    if not delete_keys or not primary_key:
        return

    # Convert numpy types to Python native types
    clean_keys = [convert_numpy_types(key) for key in delete_keys]

    # Normalize schema: empty string -> None
    schema_normalized = normalize_schema(schema)

    # Normalize primary key values format for fullmetalalchemy
    # For single PK: list of values [val1, val2, ...]
    # For composite PK: list of tuples [(pk1a, pk2a), (pk1b, pk2b), ...]
    if isinstance(primary_key, str):
        # Single column PK - values should be list of single values
        primary_key_values = clean_keys
    else:
        # Composite PK - ensure all values are tuples
        primary_key_values = []
        for key_val in clean_keys:
            if isinstance(key_val, (tuple, list)):
                # Already tuple/list - use as-is
                if len(key_val) == len(primary_key):
                    primary_key_values.append(tuple(key_val))
            else:
                # Should not happen for composite PK, but handle gracefully
                continue

    # Use fullmetalalchemy connection-based delete
    fac.delete_records_connection(
        table=table_name,
        connection=connection,
        primary_key_values=primary_key_values,
        schema=schema_normalized,
    )


def _execute_updates(
    connection: Any,
    table_name: str,
    update_records: list[dict[str, Any]],
    primary_key: str | list[str] | None,
    schema: str | None = None,
) -> None:
    """
    Execute update operations.

    Args:
        connection: SQLAlchemy connection within a transaction
        table_name: Name of the table
        update_records: List of records to update (must include primary key)
        primary_key: Name of the primary key column(s) - single string or list for composite
        schema: Optional schema name
    """
    if not update_records or not primary_key:
        return

    # Normalize schema: empty string -> None
    schema_normalized = normalize_schema(schema)

    # Convert numpy types in records
    clean_records = [convert_records_list([record])[0] for record in update_records]

    # Normalize primary key to list format for match_column_names
    if isinstance(primary_key, str):
        match_columns = [primary_key]
    else:
        match_columns = list(primary_key)

    # Use fullmetalalchemy connection-based update
    # match_column_names uses these columns to match records (supports composite PKs)
    # Records should include PK values for matching
    fac.update_records_connection(
        table=table_name,
        records=clean_records,
        connection=connection,
        match_column_names=match_columns,
        schema=schema_normalized,
    )


def _execute_inserts(
    connection: Any,
    table_name: str,
    insert_records: list[dict[str, Any]],
    schema: str | None = None,
) -> None:
    """
    Execute insert operations.

    Args:
        connection: SQLAlchemy connection within a transaction
        table_name: Name of the table
        insert_records: List of records to insert
        schema: Optional schema name
    """
    if not insert_records:
        return

    # Convert numpy types and Timestamps to Python native types in records
    # Timestamps need to be converted to datetime or string for SQLite compatibility
    clean_records = []
    for record in insert_records:
        clean_record = {}
        for key, value in record.items():
            # Convert Timestamp to datetime for database compatibility
            if isinstance(value, pd.Timestamp):
                clean_record[key] = value.to_pydatetime()
            else:
                clean_record[key] = convert_numpy_types(value)
        clean_records.append(clean_record)

    if not clean_records:
        return

    # Normalize schema: empty string -> None
    schema_normalized = normalize_schema(schema)

    # Use fullmetalalchemy connection-based insert
    fac.insert_records_connection(
        table=table_name, records=clean_records, connection=connection, schema=schema_normalized
    )


def _pandas_dtype_to_python_type(dtype: Any) -> type:
    """
    Convert pandas dtype to Python type for transmutation.

    This is a wrapper around the utils function for internal use.

    Args:
        dtype: Pandas dtype

    Returns:
        Python type (int, float, str, bool)
    """
    return pandas_dtype_to_python_type(dtype)


def create_table_from_dataframe(
    engine: Engine,
    table_name: str,
    df: pd.DataFrame,
    primary_key: str | list[str],
    schema: str | None,
    if_exists: str = "fail",
) -> None:
    """
    Create a new table from a DataFrame.

    Args:
        engine: SQLAlchemy engine
        table_name: Name of the table to create
        df: DataFrame to create table from
        primary_key: Name of the primary key column(s) - single string or list for composite
        schema: Optional schema name
        if_exists: What to do if table exists ('fail', 'replace', 'append')
            Note: 'append' is not supported by fullmetalalchemy and will raise ValueError

    Raises:
        ValueError: If table exists and if_exists='fail', or if if_exists='append'
    """
    # Handle 'append' which is not supported by fullmetalalchemy
    if if_exists == "append":
        raise ValueError(
            "if_exists='append' is not supported. Use 'fail' or 'replace', "
            "or manually insert data after table creation."
        )

    # Map 'fail' to 'error' (fullmetalalchemy uses 'error')
    if_exists_mapped = "error" if if_exists == "fail" else if_exists

    # Convert Timestamps to datetime strings for SQLite compatibility
    # SQLite doesn't natively support Timestamp objects
    df = df.copy()
    for col in df.columns:
        if df[col].dtype == "datetime64[ns]":
            # Convert to string format that SQLite can handle
            # fullmetalalchemy will convert strings to appropriate datetime type
            df[col] = pd.to_datetime(df[col]).dt.strftime("%Y-%m-%d %H:%M:%S")

    # Prepare primary key (handles index naming and validation)
    from pandalchemy.utils import prepare_primary_key_for_table_creation

    df_prepared = prepare_primary_key_for_table_creation(df, primary_key)

    # Convert DataFrame to ensure primary key is a column
    df_to_write = extract_primary_key_column(df_prepared, primary_key)

    # Ensure Timestamps are converted in the final DataFrame too
    for col in df_to_write.columns:
        if df_to_write[col].dtype == "datetime64[ns]":
            # Convert to string format that SQLite can handle
            df_to_write[col] = pd.to_datetime(df_to_write[col]).dt.strftime("%Y-%m-%d %H:%M:%S")
        # Also check if values are Timestamp objects (might happen if dtype is object)
        elif (
            df_to_write[col].dtype == "object"
            and len(df_to_write[col]) > 0
            and isinstance(df_to_write[col].iloc[0], pd.Timestamp)
        ):
            df_to_write[col] = pd.to_datetime(df_to_write[col]).dt.strftime("%Y-%m-%d %H:%M:%S")

    # Normalize schema: empty string -> None (SQLite doesn't handle empty schema strings)
    schema_normalized = normalize_schema(schema)

    # Use fullmetalalchemy to create table from DataFrame
    # This handles MySQL VARCHAR length, empty DataFrames, SQLite DateTime, etc.
    fa.create.create_table_from_dataframe(
        table_name=table_name,
        df=df_to_write,
        primary_key=primary_key,
        engine=engine,
        schema=schema_normalized,
        if_exists=if_exists_mapped,
    )


def table_exists(engine: Engine, table_name: str, schema: str | None = None) -> bool:
    """
    Check if a table exists in the database.

    Args:
        engine: SQLAlchemy engine
        table_name: Name of the table
        schema: Optional schema name

    Returns:
        True if table exists, False otherwise
    """
    return fa.features.table_exists(table_name, engine, schema)
