import pandas as pd
from typing import Optional

from ..schema import FeatureSchema

from .._core import get_logger


_LOGGER = get_logger("Data Exploration: Schema Ops")


__all__ = [
    "finalize_feature_schema",
    "apply_feature_schema",
    "reconstruct_from_schema",
]


def finalize_feature_schema(
    df_features: pd.DataFrame,
    categorical_mappings: Optional[dict[str, dict[str, int]]]
) -> FeatureSchema:
    """
    Analyzes the final features DataFrame to create a definitive schema.

    This function is the "single source of truth" for column order
    and type (categorical vs. continuous) for the entire ML pipeline.

    It should be called at the end of the feature engineering process.

    Args:
        df_features (pd.DataFrame):
            The final, processed DataFrame containing *only* feature columns
            in the exact order they will be fed to the model.
        categorical_mappings (Dict[str, Dict[str, int]] | None):
            The mappings dictionary generated by
            `encode_categorical_features`. Can be None if no
            categorical features exist.

    Returns:
        FeatureSchema: A NamedTuple containing all necessary metadata for the pipeline.
    """
    feature_names: list[str] = df_features.columns.to_list()
    
    # Intermediate lists for building
    continuous_feature_names_list: list[str] = []
    categorical_feature_names_list: list[str] = []
    categorical_index_map_dict: dict[int, int] = {}

    # _LOGGER.info("Finalizing feature schema...")

    if categorical_mappings:
        # --- Categorical features are present ---
        categorical_names_set = set(categorical_mappings.keys())
        
        for index, name in enumerate(feature_names):
            if name in categorical_names_set:
                # This is a categorical feature
                cardinality = len(categorical_mappings[name])
                categorical_index_map_dict[index] = cardinality
                categorical_feature_names_list.append(name)
            else:
                # This is a continuous feature
                continuous_feature_names_list.append(name)
        
        # Use the populated dict, or None if it's empty
        final_index_map = categorical_index_map_dict if categorical_index_map_dict else None
    
    else:
        # --- No categorical features ---
        _LOGGER.info("No categorical mappings provided. Treating all features as continuous.")
        continuous_feature_names_list = list(feature_names)
        # categorical_feature_names_list remains empty
        # categorical_index_map_dict remains empty
        final_index_map = None # Explicitly set to None to match Optional type

    _LOGGER.info(f"Schema created: {len(continuous_feature_names_list)} continuous, {len(categorical_feature_names_list)} categorical.")
    
    # Create the final immutable instance
    schema_instance = FeatureSchema(
        feature_names=tuple(feature_names),
        continuous_feature_names=tuple(continuous_feature_names_list),
        categorical_feature_names=tuple(categorical_feature_names_list),
        categorical_index_map=final_index_map,
        categorical_mappings=categorical_mappings
    )
    
    return schema_instance


def apply_feature_schema(
    df: pd.DataFrame,
    schema: FeatureSchema,
    targets: Optional[list[str]] = None,
    unknown_value: int = 99999,
    verbose: int = 3
) -> pd.DataFrame:
    """
    Aligns the input DataFrame with the provided FeatureSchema.

    This function aligns data for inference/fine-tuning by enforcing the schema's
    structure and encoding.

    Args:
        df (pd.DataFrame): The input DataFrame.
        schema (FeatureSchema): The schema defining feature names, types, and mappings.
        targets (list[str] | None): Optional list of target column names.
        unknown_value (int): Integer value to assign to unknown categorical levels.
                             Defaults to 99999 to avoid collision with existing categories.
        verbose (int): Verbosity level for logging. Higher values produce more detailed logs.

    Returns:
        pd.DataFrame: A new DataFrame with the exact column order and encoding defined by the schema.

    Raises:
        ValueError: If any required feature or target column is missing.
    """
    # 1. Setup
    df_processed = df.copy()
    targets = targets if targets is not None else []
    
    # 2. Validation: Strict Column Presence
    missing_features = [col for col in schema.feature_names if col not in df_processed.columns]
    if missing_features:
        _LOGGER.error(f"Schema Mismatch: Missing required features: {missing_features}")
        raise ValueError()
    
    # target columns should not be part of feature columns
    if targets:
        overlapping_columns = set(schema.feature_names).intersection(set(targets))
        if overlapping_columns:
            _LOGGER.error(f"Schema Mismatch: Target columns overlap with feature columns: {overlapping_columns}")
            raise ValueError()
        
        # targets were provided, check their presence
        missing_targets = [col for col in targets if col not in df_processed.columns]
        if missing_targets:
            _LOGGER.error(f"Target Mismatch: Missing target columns: {missing_targets}")
            raise ValueError()

    # 3. Apply Categorical Encoding
    if schema.categorical_feature_names and schema.categorical_mappings:
        for col_name in schema.categorical_feature_names:
            # Should never happen due to schema construction, but double-check and raise
            if col_name not in schema.categorical_mappings:
                _LOGGER.error(f"Schema Inconsistency: No mapping found for categorical feature '{col_name}'.")
                raise ValueError()

            mapping = schema.categorical_mappings[col_name]
            
            # Apply mapping (unknowns become NaN)
            df_processed[col_name] = df_processed[col_name].astype(str).map(mapping)
            
            # Handle Unknown Categories
            if df_processed[col_name].isnull().any():
                n_missing = df_processed[col_name].isnull().sum()
                if verbose >= 1:
                    _LOGGER.warning(f"Feature '{col_name}': Found {n_missing} unknown categories. Mapping to {unknown_value}.")
                
                # Fill unknowns with the specified integer
                df_processed[col_name] = df_processed[col_name].fillna(unknown_value)
            
            df_processed[col_name] = df_processed[col_name].astype(int)

    # 4. Reorder and Filter
    final_column_order = list(schema.feature_names) + targets
    
    extra_cols = set(df_processed.columns) - set(final_column_order)
    if extra_cols:
        if verbose >= 1:
            _LOGGER.warning(f"Dropping {len(extra_cols)} extra columns not present in schema: {extra_cols}")

    df_final = df_processed[final_column_order]
    
    if verbose >= 2:
        _LOGGER.info(f"Schema applied successfully. Final shape: {df_final.shape}")
    
    # df_final should be a dataframe
    if isinstance(df_final, pd.Series):
        df_final = df_final.to_frame()

    return df_final


def reconstruct_from_schema(
    df: pd.DataFrame,
    schema: FeatureSchema,
    targets: Optional[list[str]] = None,
    verbose: int = 3
) -> pd.DataFrame:
    """
    Reverses the schema application to make data human-readable.

    This function decodes categorical features back to their string representations
    using the schema's mappings. It strictly enforces the schema structure,
    ignoring extra columns (unless they are specified as targets).

    Args:
        df (pd.DataFrame): The input DataFrame containing encoded features.
        schema (FeatureSchema): The schema defining feature names and reverse mappings.
        targets (list[str] | None): Optional list of target column names to preserve. These are not decoded and kept in the order specified here.
        verbose (int): Verbosity level for logging info about the process.

    Returns:
        pd.DataFrame: A new DataFrame with the exact column order (features + targets),
                      with categorical features decoded to strings.

    Raises:
        ValueError: If any required feature or target column is missing.
    """
    # 1. Setup
    df_decoded = df.copy()
    targets = targets if targets is not None else []

    # 2. Validation: Strict Column Presence
    # Check Features
    missing_features = [col for col in schema.feature_names if col not in df_decoded.columns]
    if missing_features:
        _LOGGER.error(f"Schema Reconstruction Mismatch: Missing required features: {missing_features}")
        raise ValueError()
    
    # Check Targets
    if targets:
        missing_targets = [col for col in targets if col not in df_decoded.columns]
        if missing_targets:
            _LOGGER.error(f"Schema Reconstruction Mismatch: Missing required targets: {missing_targets}")
            raise ValueError()

    # 3. Reorder and Filter (Drop extra columns early)
    # The valid columns are Features + Targets
    valid_columns = list(schema.feature_names) + targets
    
    extra_cols = set(df_decoded.columns) - set(valid_columns)
    if extra_cols:
        if verbose >= 1:
            _LOGGER.warning(f"Dropping extra columns not present in schema or targets: {extra_cols}")

    # Enforce order: Features first, then Targets
    df_decoded = df_decoded[valid_columns]

    # 4. Reverse Categorical Encoding
    if schema.categorical_feature_names and schema.categorical_mappings:
        for col_name in schema.categorical_feature_names:
            if col_name not in schema.categorical_mappings:
                continue

            forward_mapping = schema.categorical_mappings[col_name]
            # Create reverse map: {int: str}
            reverse_mapping = {v: k for k, v in forward_mapping.items()}
            
            # --- SAFE TYPE CASTING ---
            # Ensure values are Integers before mapping (handle 5.0 vs 5).
            try:
                if pd.api.types.is_numeric_dtype(df_decoded[col_name]):
                    df_decoded[col_name] = df_decoded[col_name].astype("Int64")
            except (TypeError, ValueError):
                # casted to NaN later during mapping
                pass
            # -------------------------

            # Check for unknown codes before mapping
            if verbose >= 1:
                unique_codes = df_decoded[col_name].dropna().unique()
                unknown_codes = [code for code in unique_codes if code not in reverse_mapping]
                if unknown_codes:
                    _LOGGER.warning(f"Feature '{col_name}': Found unknown encoded values {unknown_codes}. These will be mapped to NaN.")

            # Apply reverse mapping
            df_decoded[col_name] = df_decoded[col_name].map(reverse_mapping)
            
    if verbose >= 2:
        _LOGGER.info(f"Schema reconstruction successful. Final shape: {df_decoded.shape}")
    
    return df_decoded

