"""
Enhanced Structural Fixers - Level 2.5

Structural fixes beyond basic value cleaning:
- Column reconstruction
- Schema merging
- Type drift resolution  
- Column normalization
- JSON flattening
"""

from typing import List, Dict, Optional
import pandas as pd
import json
import re


class StructuralFixer:
    """Advanced structural fixes."""
    
    def reconstruct_missing_columns(
        self,
        df: pd.DataFrame,
        expected_columns: List[str],
        fill_strategy: str = 'na'
    ) -> pd.DataFrame:
        """
        Reconstruct missing columns.
        
        Args:
            df: DataFrame to fix
            expected_columns: List of expected column names
            fill_strategy: 'na', 'zero', 'empty', or 'infer'
        """
        df_copy = df.copy()
        
        for col in expected_columns:
            if col not in df_copy.columns:
                if fill_strategy == 'na':
                    df_copy[col] = pd.NA
                elif fill_strategy == 'zero':
                    df_copy[col] = 0
                elif fill_strategy == 'empty':
                    df_copy[col] = ''
                else:  # infer
                    df_copy[col] = pd.NA
        
        return df_copy
    
    def merge_conflicting_schemas(
        self,
        dataframes: Dict[str, pd.DataFrame],
        strategy: str = 'union'
    ) -> pd.DataFrame:
        """
        Merge files with different schemas.
        
        Args:
            dataframes: Dict of filename -> DataFrame
            strategy: 'union' (all columns) or 'intersection' (common only)
        """
        if not dataframes:
            return pd.DataFrame()
        
        if strategy == 'union':
            # Get all unique columns
            all_columns = list(dict.fromkeys([  # Preserve order
                col
                for df in dataframes.values()
                for col in df.columns
            ]))
            
            # Add missing columns to each dataframe
            aligned = []
            for df in dataframes.values():
                df_copy = df.copy()
                for col in all_columns:
                    if col not in df_copy.columns:
                        df_copy[col] = pd.NA
                df_copy = df_copy[all_columns]  # Reorder
                aligned.append(df_copy)
            
            # Concatenate
            return pd.concat(aligned, ignore_index=True)
        
        else:  # intersection
            # Get common columns
            common_cols = set.intersection(*[set(df.columns) for df in dataframes.values()])
            common_cols = list(common_cols)
            
            # Select common columns only
            aligned = [df[common_cols] for df in dataframes.values()]
            
            return pd.concat(aligned, ignore_index=True)
    
    def resolve_type_drift(
        self,
        df: pd.DataFrame,
        column: str,
        target_type: str = 'string'
    ) -> pd.DataFrame:
        """
        Resolve type drift by converting to safe type.
        
        Args:
            df: DataFrame to fix
            column: Column name
            target_type: Target type ('string', 'float', 'int')
        """
        df_copy = df.copy()
        
        if column not in df_copy.columns:
            return df_copy
        
        try:
            if target_type == 'string':
                df_copy[column] = df_copy[column].astype(str)
            elif target_type == 'float':
                df_copy[column] = pd.to_numeric(df_copy[column], errors='coerce')
            elif target_type == 'int':
                df_copy[column] = pd.to_numeric(df_copy[column], errors='coerce').astype('Int64')
        except Exception:
            # Fallback to string
            df_copy[column] = df_copy[column].astype(str)
        
        return df_copy
    
    def standardize_column_order(
        self,
        dataframes: Dict[str, pd.DataFrame],
        column_order: Optional[List[str]] = None
    ) -> Dict[str, pd.DataFrame]:
        """
        Standardize column order across all files.
        
        Args:
            dataframes: Dict of filename -> DataFrame
            column_order: Desired column order (inferred if None)
        """
        if not dataframes:
            return {}
        
        if column_order is None:
            # Use order from first file with most columns
            column_order = max(
                dataframes.values(),
                key=lambda df: len(df.columns)
            ).columns.tolist()
        
        standardized = {}
        for name, df in dataframes.items():
            # Reorder existing columns
            existing_cols = [c for c in column_order if c in df.columns]
            standardized[name] = df[existing_cols]
        
        return standardized
    
    def normalize_column_names(
        self,
        df: pd.DataFrame,
        style: str = 'snake_case'
    ) -> pd.DataFrame:
        """
        Normalize column names to consistent format.
        
        Args:
            df: DataFrame to normalize
            style: 'snake_case', 'camelCase', or 'lowercase'
        """
        df_copy = df.copy()
        
        if style == 'snake_case':
            new_columns = {}
            for col in df_copy.columns:
                # Convert to snake_case
                normalized = col.strip()
                normalized = re.sub(r'[^\w\s]', '', normalized)  # Remove special chars
                normalized = re.sub(r'\s+', '_', normalized)  # Spaces to underscores
                normalized = re.sub(r'([a-z0-9])([A-Z])', r'\1_\2', normalized)  # camelCase -> snake_case
                normalized = normalized.lower()
                normalized = re.sub(r'_+', '_', normalized)  # Multiple underscores -> one
                normalized = normalized.strip('_')
                new_columns[col] = normalized
            
            df_copy = df_copy.rename(columns=new_columns)
        
        elif style == 'camelCase':
            new_columns = {}
            for col in df_copy.columns:
                words = re.split(r'[^\w]+', col.strip())
                if words:
                    normalized = words[0].lower() + ''.join(w.capitalize() for w in words[1:])
                    new_columns[col] = normalized
            
            df_copy = df_copy.rename(columns=new_columns)
        
        else:  # lowercase
            df_copy.columns = [col.lower().strip() for col in df_copy.columns]
        
        return df_copy
    
    def flatten_nested_json(
        self,
        df: pd.DataFrame,
        json_column: str,
        prefix: Optional[str] = None
    ) -> pd.DataFrame:
        """
        Flatten JSON objects into separate columns.
        
        Args:
            df: DataFrame to flatten
            json_column: Column containing JSON strings
            prefix: Prefix for new columns (uses json_column if None)
        """
        if json_column not in df.columns:
            return df
        
        if prefix is None:
            prefix = json_column
        
        df_copy = df.copy()
        
        # Parse JSON and flatten
        json_data = []
        for val in df_copy[json_column]:
            if pd.isna(val):
                json_data.append({})
            elif isinstance(val, str):
                try:
                    parsed = json.loads(val)
                    if isinstance(parsed, dict):
                        json_data.append(parsed)
                    else:
                        json_data.append({})
                except:
                    json_data.append({})
            else:
                json_data.append({})
        
        # Convert to DataFrame
        if json_data:
            flattened = pd.json_normalize(json_data)
            # Add prefix
            flattened.columns = [f"{prefix}_{col}" for col in flattened.columns]
            
            # Drop original JSON column and concatenate flattened
            df_copy = df_copy.drop(columns=[json_column])
            df_copy = pd.concat([df_copy, flattened], axis=1)
        
        return df_copy
    
    def remove_duplicate_columns(
        self,
        df: pd.DataFrame,
        fuzzy_threshold: float = 0.9
    ) -> pd.DataFrame:
        """
        Remove duplicate columns (exact or fuzzy match).
        
        Args:
            df: DataFrame to clean
            fuzzy_threshold: Similarity threshold for fuzzy matching
        """
        from difflib import SequenceMatcher
        
        df_copy = df.copy()
        columns_to_drop = set()
        
        cols = list(df_copy.columns)
        for i in range(len(cols)):
            if cols[i] in columns_to_drop:
                continue
            
            for j in range(i + 1, len(cols)):
                if cols[j] in columns_to_drop:
                    continue
                
                # Check similarity
                similarity = SequenceMatcher(None, cols[i].lower(), cols[j].lower()).ratio()
                
                if similarity >= fuzzy_threshold:
                    # Keep first, drop second
                    columns_to_drop.add(cols[j])
        
        if columns_to_drop:
            df_copy = df_copy.drop(columns=list(columns_to_drop))
        
        return df_copy
