"""
Polars-based data cleaning engine (primary, fast, in-memory).
"""
import polars as pl
from pathlib import Path
from typing import Dict, Any, List
import structlog

logger = structlog.get_logger(__name__)


class CleanResult:
    """Result of a cleaning operation."""
    
    def __init__(self, df: pl.DataFrame, issues_found: int, fixes_applied: int, quality_score: float):
        self.df = df
        self.issues_found = issues_found
        self.fixes_applied = fixes_applied
        self.quality_score = quality_score
    
    def save(self, output_path: str):
        """Save cleaned DataFrame to CSV."""
        self.df.write_csv(output_path)
    
    def generate_report(self, report_path: str):
        """Generate HTML quality report."""
        # Placeholder - will use existing report generator
        html = f"""
        <html>
        <head><title>Data Quality Report</title></head>
        <body>
            <h1>Quality Score: {self.quality_score:.1f}%</h1>
            <p>Issues found: {self.issues_found}</p>
            <p>Fixes applied: {self.fixes_applied}</p>
        </body>
        </html>
        """
        Path(report_path).write_text(html)
    
    def create_bundle(self, bundle_path: str):
        """Create ZIP bundle with cleaned data, report, and schema."""
        import zipfile
        import json
        
        with zipfile.ZipFile(bundle_path, 'w') as zf:
            # Add cleaned CSV
            temp_csv = "temp_cleaned.csv"
            self.save(temp_csv)
            zf.write(temp_csv, "cleaned.csv")
            Path(temp_csv).unlink()
            
            # Add report
            temp_report = "temp_report.html"
            self.generate_report(temp_report)
            zf.write(temp_report, "report.html")
            Path(temp_report).unlink()
            
            # Add schema
            schema = {col: str(dtype) for col, dtype in zip(self.df.columns, self.df.dtypes)}
            zf.writestr("schema.json", json.dumps(schema, indent=2))


class PolarsEngine:
    """Fast in-memory data cleaning with Polars."""
    
    def __init__(self):
        self.name = "Polars"
    
    def clean(self, file_path: str, conservative: bool = False) -> CleanResult:
        """
        Clean a CSV file using Polars.
        
        Args:
            file_path: Path to CSV file
            conservative: If True, make minimal changes
            
        Returns:
            CleanResult with cleaned DataFrame and stats
        """
        logger.info("polars_engine_start", file=file_path, conservative=conservative)
        
        try:
            # Read CSV
            df = pl.read_csv(file_path, infer_schema_length=10000)
            
            original_rows = len(df)
            issues_found = 0
            fixes_applied = 0
            
            # Detect and fix issues
            # 1. Remove duplicate rows
            duplicates = len(df) - df.n_unique()
            if duplicates > 0:
                issues_found += duplicates
                df = df.unique()
                fixes_applied += duplicates
                logger.info("removed_duplicates", count=duplicates)
            
            # 2. Handle missing values
            for col in df.columns:
                nulls = df[col].null_count()
                if nulls > 0:
                    issues_found += nulls
                    
                    # Fill strategy based on dtype
                    if df[col].dtype == pl.Int64 or df[col].dtype == pl.Float64:
                        # Fill numeric with median
                        median_val = df[col].median()
                        if median_val is not None:
                            df = df.with_columns(df[col].fill_null(median_val))
                            fixes_applied += nulls
                    else:
                        # Fill categorical with mode or "UNKNOWN"
                        if not conservative:
                            df = df.with_columns(df[col].fill_null("UNKNOWN"))
                            fixes_applied += nulls
            
            # 3. Trim whitespace from strings
            for col in df.columns:
                if df[col].dtype == pl.Utf8:
                    df = df.with_columns(df[col].str.strip_chars())
            
            # Calculate quality score
            cleaned_rows = len(df)
            quality_score = ((cleaned_rows / original_rows) * 100) if original_rows > 0 else 100.0
            
            logger.info("polars_engine_complete", 
                       issues=issues_found, 
                       fixes=fixes_applied, 
                       quality=quality_score)
            
            return CleanResult(
                df=df,
                issues_found=issues_found,
                fixes_applied=fixes_applied,
                quality_score=quality_score
            )
            
        except Exception as e:
            logger.error("polars_engine_failed", error=str(e))
            raise
