"""
Schema Reconciler - Multi-file schema reconciliation and harmonization.

This is the CORE of Level 2.5 - structural healing at scale.
"""

from typing import List, Dict, Set, Tuple, Optional
from dataclasses import dataclass
import pandas as pd
from difflib import SequenceMatcher
from collections import Counter


@dataclass
class ColumnMapping:
    """Mapping between column names across files."""
    canonical_name: str
    variations: List[str]
    confidence: float
    source_files: List[str]


@dataclass  
class SchemaDrift:
    """Schema drift detected across files."""
    missing_columns: Dict[str, List[str]]  # column -> files where it's missing
    extra_columns: Dict[str, List[str]]     # column -> files where it appears
    type_drifts: Dict[str, Dict[str, str]]  # column -> {file: type}
    column_order_drifts: List[str]          # files with different column orders


@dataclass
class UnifiedSchema:
    """Unified schema across all files."""
    columns: List[str]  # Canonical column order
    column_types: Dict[str, str]
    nullable: Dict[str, bool]
    mappings: List[ColumnMapping]
    confidence_score: float


class SchemaReconciler:
    """
    Reconciles schemas across multiple files.
    
    Handles:
    - Fuzzy column name matching
    - Missing column detection
    - Type drift resolution
    - Column order standardization
    """
    
    def __init__(self, fuzzy_threshold: float = 0.85):
        """
        Args:
            fuzzy_threshold: Minimum similarity for fuzzy matching (0-1)
        """
        self.fuzzy_threshold = fuzzy_threshold
    
    def analyze_schemas(
        self,
        dataframes: Dict[str, pd.DataFrame]
    ) -> Tuple[SchemaDrift, UnifiedSchema]:
        """
        Analyze schemas across multiple DataFrames.
        
        Args:
            dataframes: Dict of filename -> DataFrame
            
        Returns:
            Tuple of (SchemaDrift, UnifiedSchema)
        """
        # Collect all column names from all files
        all_columns = {}
        for filename, df in dataframes.items():
            all_columns[filename] = list(df.columns)
        
        # Detect column name variations using fuzzy matching
        column_groups = self._group_similar_columns(all_columns)
        
        # Create canonical names
        mappings = self._create_column_mappings(column_groups, all_columns)
        
        # Detect drift
        drift = self._detect_drift(dataframes, mappings)
        
        # Create unified schema
        unified = self._create_unified_schema(dataframes, mappings)
        
        return drift, unified
    
    def _group_similar_columns(
        self,
        all_columns: Dict[str, List[str]]
    ) -> List[List[Tuple[str, str]]]:
        """
        Group similar column names using fuzzy matching.
        
        Returns:
            List of groups, each group is list of (filename, column_name) tuples
        """
        # Flatten all columns with their source files
        column_pairs = []
        for filename, columns in all_columns.items():
            for col in columns:
                column_pairs.append((filename, col))
        
        # Group similar columns
        groups = []
        used = set()
        
        for i, (file1, col1) in enumerate(column_pairs):
            if (file1, col1) in used:
                continue
            
            group = [(file1, col1)]
            used.add((file1, col1))
            
            # Find similar columns
            for j, (file2, col2) in enumerate(column_pairs):
                if i == j or (file2, col2) in used:
                    continue
                
                similarity = self._column_similarity(col1, col2)
                if similarity >= self.fuzzy_threshold:
                    group.append((file2, col2))
                    used.add((file2, col2))
            
            groups.append(group)
        
        return groups
    
    def _column_similarity(self, col1: str, col2: str) -> float:
        """
        Calculate similarity between two column names.
        
        Handles:
        - Case differences: "Email" vs "email"
        - Separator differences: "user_id" vs "userId" vs "user id"
        - Minor typos: "adress" vs "address"
        """
        # Normalize
        norm1 = self._normalize_column_name(col1)
        norm2 = self._normalize_column_name(col2)
        
        # Exact match after normalization
        if norm1 == norm2:
            return 1.0
        
        # Fuzzy match
        return SequenceMatcher(None, norm1, norm2).ratio()
    
    def _normalize_column_name(self, col: str) -> str:
        """Normalize column name for comparison."""
        import re
        # Lowercase
        col = col.lower()
        # Remove special characters
        col = re.sub(r'[^a-z0-9]+', '_', col)
        # Remove leading/trailing underscores
        col = col.strip('_')
        return col
    
    def _create_column_mappings(
        self,
        column_groups: List[List[Tuple[str, str]]],
        all_columns: Dict[str, List[str]]
    ) -> List[ColumnMapping]:
        """Create canonical column mappings."""
        mappings = []
        
        for group in column_groups:
            # Get all unique column names in this group
            col_names = [col for _, col in group]
            
            # Choose canonical name (most common, or alphabetically first)
            canonical = Counter(col_names).most_common(1)[0][0]
            
            # Calculate confidence based on group size vs total files
            confidence = len(set(f for f, _ in group)) / len(all_columns)
            
            mapping = ColumnMapping(
                canonical_name=canonical,
                variations=list(set(col_names)),
                confidence=confidence,
                source_files=[f for f, _ in group]
            )
            mappings.append(mapping)
        
        return mappings
    
    def _detect_drift(
        self,
        dataframes: Dict[str, pd.DataFrame],
        mappings: List[ColumnMapping]
    ) -> SchemaDrift:
        """Detect schema drift across files."""
        # Get canonical columns
        canonical_cols = {m.canonical_name for m in mappings}
        
        # Detect missing columns
        missing = {}
        for col in canonical_cols:
            missing_in = []
            for filename, df in dataframes.items():
                # Check if this canonical column exists in this file
                mapping = next((m for m in mappings if m.canonical_name == col), None)
                if mapping:
                    has_col = any(var in df.columns for var in mapping.variations)
                    if not has_col:
                        missing_in.append(filename)
            
            if missing_in:
                missing[col] = missing_in
        
        # Detect extra columns
        extra = {}
        for filename, df in dataframes.items():
            for col in df.columns:
                # Check if this column is mapped
                is_mapped = any(col in m.variations for m in mappings)
                if not is_mapped:
                    if col not in extra:
                        extra[col] = []
                    extra[col].append(filename)
        
        # Detect type drifts
        type_drifts = {}
        for mapping in mappings:
            types_by_file = {}
            for filename, df in dataframes.items():
                # Find the column in this file
                col_in_file = next((c for c in df.columns if c in mapping.variations), None)
                if col_in_file:
                    types_by_file[filename] = str(df[col_in_file].dtype)
            
            # Check if types differ
            if len(set(types_by_file.values())) > 1:
                type_drifts[mapping.canonical_name] = types_by_file
        
        # Detect column order drift (simplified)
        order_drifts = []
        if len(dataframes) > 1:
            first_file = list(dataframes.keys())[0]
            first_order = list(dataframes[first_file].columns)
            
            for filename, df in list(dataframes.items())[1:]:
                if list(df.columns) != first_order:
                    order_drifts.append(filename)
        
        return SchemaDrift(
            missing_columns=missing,
            extra_columns=extra,
            type_drifts=type_drifts,
            column_order_drifts=order_drifts
        )
    
    def _create_unified_schema(
        self,
        dataframes: Dict[str, pd.DataFrame],
        mappings: List[ColumnMapping]
    ) -> UnifiedSchema:
        """Create unified schema from mappings."""
        # Determine canonical column order (by frequency across files)
        column_order = [m.canonical_name for m in sorted(
            mappings,
            key=lambda m: len(m.source_files),
            reverse=True
        )]
        
        # Determine column types (majority vote)
        column_types = {}
        nullable = {}
        
        for mapping in mappings:
            types = []
            nullables = []
            
            for filename, df in dataframes.items():
                col_in_file = next((c for c in df.columns if c in mapping.variations), None)
                if col_in_file:
                    types.append(str(df[col_in_file].dtype))
                    nullables.append(df[col_in_file].isna().any())
            
            if types:
                # Majority vote for type
                column_types[mapping.canonical_name] = Counter(types).most_common(1)[0][0]
                # Any nullable -> nullable
                nullable[mapping.canonical_name] = any(nullables)
        
        # Calculate overall confidence
        confidence = sum(m.confidence for m in mappings) / len(mappings) if mappings else 0.0
        
        return UnifiedSchema(
            columns=column_order,
            column_types=column_types,
            nullable=nullable,
            mappings=mappings,
            confidence_score=confidence
        )
    
    def reconcile_dataframes(
        self,
        dataframes: Dict[str, pd.DataFrame],
        unified_schema: UnifiedSchema
    ) -> Dict[str, pd.DataFrame]:
        """
        Reconcile all DataFrames to match unified schema.
        
        - Renames columns to canonical names
        - Adds missing columns (NaN fill)
        - Removes extra columns
        - Reorders columns
        """
        reconciled = {}
        
        for filename, df in dataframes.items():
            df_copy = df.copy()
            
            # Step 1: Rename columns to canonical names
            rename_map = {}
            for mapping in unified_schema.mappings:
                for variation in mapping.variations:
                    if variation in df_copy.columns:
                        rename_map[variation] = mapping.canonical_name
            
            df_copy = df_copy.rename(columns=rename_map)
            
            # Step 2: Add missing columns
            for col in unified_schema.columns:
                if col not in df_copy.columns:
                    df_copy[col] = pd.NA
            
            # Step 3: Remove extra columns
            extra_cols = [c for c in df_copy.columns if c not in unified_schema.columns]
            if extra_cols:
                df_copy = df_copy.drop(columns=extra_cols)
            
            # Step 4: Reorder columns
            df_copy = df_copy[unified_schema.columns]
            
            reconciled[filename] = df_copy
        
        return reconciled
