"""
Structural Engine - Level 2.5 structural healing.

Features:
- Column alignment across entire dataset
- Multi-file header harmonization
- Row-length mismatch repair
- Corrupt block isolation
- JSON flattening
- Canonical column ordering
- Missing column reconstruction
- Deterministic schema merge
"""

import polars as pl
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass


@dataclass
class StructuralIssue:
    """Structural issue detected."""
    issue_type: str
    severity: str
    description: str
    affected_rows: List[int]
    fixable: bool


class StructuralEngine:
    """Structural healing engine using Polars."""
    
    def align_columns(
        self,
        dfs: Dict[str, pl.DataFrame],
        canonical_order: Optional[List[str]] = None
    ) -> Dict[str, pl.DataFrame]:
        """
        Align columns across entire dataset.
        
        Ensures all DataFrames have same columns in same order.
        """
        # Determine canonical order
        if canonical_order is None:
            # Use union of all columns
            all_cols = set()
            for df in dfs.values():
                all_cols.update(df.columns)
            canonical_order = sorted(all_cols)
        
        aligned = {}
        for name, df in dfs.items():
            # Add missing columns
            for col in canonical_order:
                if col not in df.columns:
                    df = df.with_columns(pl.lit(None).alias(col))
            
            # Reorder
            df = df.select(canonical_order)
            aligned[name] = df
        
        return aligned
    
    def harmonize_headers(
        self,
        dfs: Dict[str, pl.DataFrame],
        fuzzy_threshold: float = 0.85
    ) -> Dict[str, pl.DataFrame]:
        """
        Harmonize headers across multiple files using fuzzy matching.
        
        Similar to PolarsSchemaReconciler but focused on headers only.
        """
        from difflib import SequenceMatcher
        
        # Group similar headers
        all_headers = {}
        for name, df in dfs.items():
            all_headers[name] = list(df.columns)
        
        # Find canonical names
        header_groups = {}
        for name, headers in all_headers.items():
            for header in headers:
                matched = False
                for canonical in header_groups.keys():
                    similarity = SequenceMatcher(None, header.lower(), canonical.lower()).ratio()
                    if similarity >= fuzzy_threshold:
                        header_groups[canonical].append((name, header))
                        matched = True
                        break
                
                if not matched:
                    header_groups[header] = [(name, header)]
        
        # Rename to canonical
        harmonized = {}
        for name, df in dfs.items():
            rename_map = {}
            for canonical, variants in header_groups.items():
                for file_name, header in variants:
                    if file_name == name and header in df.columns:
                        rename_map[header] = canonical
            
            harmonized[name] = df.rename(rename_map)
        
        return harmonized
    
    def detect_row_length_mismatch(
        self,
        df: pl.DataFrame,
        expected_columns: int
    ) -> List[StructuralIssue]:
        """
        Detect rows with wrong number of values.
        
        Note: Polars handles this during CSV reading, but we can
        detect if columns were auto-padded.
        """
        issues = []
        
        # Check for columns that are entirely null
        for col in df.columns:
            if df[col].null_count() == df.height:
                issues.append(StructuralIssue(
                    issue_type='EMPTY_COLUMN',
                    severity='medium',
                    description=f'Column {col} is entirely null',
                    affected_rows=[],
                    fixable=True
                ))
        
        return issues
    
    def repair_row_length(
        self,
        df: pl.DataFrame,
        expected_columns: int,
        strategy: str = 'truncate'
    ) -> pl.DataFrame:
        """
        Repair rows with mismatched lengths.
        
        Strategies:
        - 'truncate': Remove extra columns
        - 'pad': Add missing columns with None
        """
        current_cols = len(df.columns)
        
        if current_cols > expected_columns and strategy == 'truncate':
            # Keep only first N columns
            df = df.select(df.columns[:expected_columns])
        
        elif current_cols < expected_columns and strategy == 'pad':
            # Add missing columns
            for i in range(current_cols, expected_columns):
                df = df.with_columns(pl.lit(None).alias(f'column_{i}'))
        
        return df
    
    def isolate_corrupt_blocks(
        self,
        df: pl.DataFrame,
        max_consecutive_nulls: int = 3
    ) -> Tuple[pl.DataFrame, pl.DataFrame]:
        """
        Isolate rows that appear to be corrupt blocks.
        
        Returns: (clean_df, corrupt_df)
        """
        # Count nulls per row
        null_counts = df.select(
            pl.sum_horizontal([pl.col(c).is_null() for c in df.columns]).alias('null_count')
        )
        
        # Mark rows with too many nulls as corrupt
        threshold = len(df.columns) * 0.7  # 70% nulls
        
        clean_mask = null_counts['null_count'] < threshold
        
        clean_df = df.filter(clean_mask)
        corrupt_df = df.filter(~clean_mask)
        
        return clean_df, corrupt_df
    
    def flatten_json_columns(
        self,
        df: pl.DataFrame,
        json_columns: Optional[List[str]] = None
    ) -> pl.DataFrame:
        """
        Flatten JSON columns into separate columns.
        
        If json_columns is None, auto-detect.
        """
        import json as json_lib
        
        if json_columns is None:
            # Auto-detect JSON columns
            json_columns = []
            for col in df.columns:
                if df[col].dtype == pl.Utf8:
                    # Check if first non-null value is JSON
                    sample = df[col].drop_nulls().limit(1)
                    if sample.height > 0:
                        try:
                            json_lib.loads(sample[0])
                            json_columns.append(col)
                        except:
                            pass
        
        # Flatten each JSON column
        for col in json_columns:
            if col not in df.columns:
                continue
            
            try:
                # Extract JSON and create struct
                df = df.with_columns(
                    pl.col(col).str.json_decode().alias(f'{col}_struct')
                )
                
                # Unnest struct
                struct_col = f'{col}_struct'
                if struct_col in df.columns:
                    df = df.unnest(struct_col)
                    df = df.drop(col)  # Drop original JSON column
            except Exception:
                # If flattening fails, keep original
                pass
        
        return df
    
    def create_canonical_order(
        self,
        dfs: Dict[str, pl.DataFrame],
        strategy: str = 'frequency'
    ) -> List[str]:
        """
        Create canonical column ordering.
        
        Strategies:
        - 'frequency': Most common order across files
        - 'alphabetical': Sort alphabetically
        - 'first': Use first file's order
        """
        if strategy == 'alphabetical':
            all_cols = set()
            for df in dfs.values():
                all_cols.update(df.columns)
            return sorted(all_cols)
        
        elif strategy == 'first':
            return list(list(dfs.values())[0].columns)
        
        else:  # frequency
            # Find most common order
            from collections import Counter
            
            orders = [tuple(df.columns) for df in dfs.values()]
            most_common = Counter(orders).most_common(1)
            
            if most_common:
                return list(most_common[0][0])
            else:
                return list(list(dfs.values())[0].columns)
    
    def reconstruct_missing_columns(
        self,
        df: pl.DataFrame,
        schema: Dict[str, str],
        fill_strategy: str = 'null'
    ) -> pl.DataFrame:
        """
        Reconstruct missing columns based on schema.
        
        Fill strategies:
        - 'null': Fill with None
        - 'zero': Fill with 0 for numeric, empty string for text
        - 'infer': Try to infer from other columns
        """
        for col_name, dtype in schema.items():
            if col_name not in df.columns:
                if fill_strategy == 'null':
                    df = df.with_columns(pl.lit(None).alias(col_name))
                elif fill_strategy == 'zero':
                    if 'int' in dtype.lower() or 'float' in dtype.lower():
                        df = df.with_columns(pl.lit(0).alias(col_name))
                    else:
                        df = df.with_columns(pl.lit('').alias(col_name))
                else:  # null by default
                    df = df.with_columns(pl.lit(None).alias(col_name))
        
        return df
    
    def deterministic_schema_merge(
        self,
        schemas: List[Dict[str, str]],
        conflict_resolution: str = 'most_common'
    ) -> Dict[str, str]:
        """
        Merge multiple schemas deterministically.
        
        Conflict resolution:
        - 'most_common': Use most common type
        - 'widest': Use widest type (e.g., string over int)
        - 'first': Use first schema's type
        """
        from collections import Counter
        
        # Get all column names
        all_cols = set()
        for schema in schemas:
            all_cols.update(schema.keys())
        
        merged = {}
        
        for col in all_cols:
            types = [s[col] for s in schemas if col in s]
            
            if not types:
                continue
            
            if conflict_resolution == 'most_common':
                merged[col] = Counter(types).most_common(1)[0][0]
            elif conflict_resolution == 'widest':
                # String is widest
                if 'Utf8' in types or 'str' in types:
                    merged[col] = 'Utf8'
                elif 'Float64' in types:
                    merged[col] = 'Float64'
                else:
                    merged[col] = types[0]
            else:  # first
                merged[col] = types[0]
        
        return merged
