"""
DuckDB-based data cleaning engine (fallback, stable, out-of-core).

Supports chunked processing for files larger than RAM.
"""
import duckdb
import polars as pl
from pathlib import Path
from typing import Optional, Callable
from datashaper.engine.polars_engine import CleanResult
import structlog

logger = structlog.get_logger(__name__)


class DuckDBEngine:
    """Stable out-of-core data cleaning with DuckDB."""
    
    # Chunk size for processing large files (rows per chunk)
    CHUNK_SIZE = 1_000_000  # 1M rows per chunk
    
    def __init__(self):
        self.name = "DuckDB"
    
    def clean(
        self, 
        file_path: str, 
        conservative: bool = False,
        progress_callback: Optional[Callable[[int, int, str], None]] = None
    ) -> CleanResult:
        """
        Clean a CSV file using DuckDB (out-of-core processing).
        
        Args:
            file_path: Path to CSV file
            conservative: If True, make minimal changes
            progress_callback: Optional callback(current_chunk, total_chunks, message)
            
        Returns:
            CleanResult with cleaned DataFrame and stats
        """
        logger.info("duckdb_engine_start", file=file_path, conservative=conservative)
        
        try:
            # Create in-memory DuckDB connection
            conn = duckdb.connect(database=':memory:')
            
            # Check file size to determine if chunking is needed
            file_size = Path(file_path).stat().st_size
            use_chunking = file_size > (1 * 1024 * 1024 * 1024)  # > 1GB
            
            if use_chunking:
                logger.info("using_chunked_processing", file_size=file_size)
                return self._clean_chunked(conn, file_path, conservative, progress_callback)
            else:
                logger.info("using_standard_processing", file_size=file_size)
                return self._clean_standard(conn, file_path, conservative)
            
        except Exception as e:
            logger.error("duckdb_engine_failed", error=str(e))
            raise
        finally:
            if 'conn' in locals():
                conn.close()
    
    def _clean_standard(
        self, 
        conn: duckdb.DuckDBPyConnection, 
        file_path: str, 
        conservative: bool
    ) -> CleanResult:
        """Standard cleaning for files < 1GB."""
        
        # Load CSV into DuckDB
        conn.execute(f"""
            CREATE TABLE raw_data AS 
            SELECT * FROM read_csv_auto('{file_path}', sample_size=-1)
        """)
        
        original_rows = conn.execute("SELECT COUNT(*) FROM raw_data").fetchone()[0]
        issues_found = 0
        fixes_applied = 0
        
        # 1. Remove duplicates
        duplicates = conn.execute("""
            SELECT COUNT(*) - COUNT(DISTINCT *) FROM raw_data
        """).fetchone()[0]
        
        if duplicates > 0:
            issues_found += duplicates
            conn.execute("CREATE TABLE data AS SELECT DISTINCT * FROM raw_data")
            fixes_applied += duplicates
            logger.info("removed_duplicates", count=duplicates)
        else:
            conn.execute("CREATE TABLE data AS SELECT * FROM raw_data")
        
        # 2. Handle missing values (count nulls)
        columns = conn.execute("PRAGMA table_info(data)").fetchall()
        for col_info in columns:
            col_name = col_info[1]  # Column name
            
            nulls = conn.execute(f"""
                SELECT COUNT(*) FROM data WHERE "{col_name}" IS NULL
            """).fetchone()[0]
            
            if nulls > 0:
                issues_found += nulls
                
                if not conservative:
                    # Simple strategy: fill with 'UNKNOWN' or 0
                    conn.execute(f"""
                        UPDATE data 
                        SET "{col_name}" = COALESCE("{col_name}", 'UNKNOWN')
                    """)
                    fixes_applied += nulls
        
        # Get cleaned data as Polars DataFrame
        df = conn.execute("SELECT * FROM data").pl()
        
        cleaned_rows = len(df)
        quality_score = ((cleaned_rows / original_rows) * 100) if original_rows > 0 else 100.0
        
        logger.info("duckdb_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
        )
    
    def _clean_chunked(
        self,
        conn: duckdb.DuckDBPyConnection,
        file_path: str,
        conservative: bool,
        progress_callback: Optional[Callable[[int, int, str], None]] = None
    ) -> CleanResult:
        """
        Chunked cleaning for large files (> 1GB).
        
        Processes file in chunks to avoid memory issues.
        """
        logger.info("chunked_processing_start", file=file_path)
        
        # Step 1: Get total row count
        if progress_callback:
            progress_callback(0, 100, "Analyzing file size...")
        
        conn.execute(f"""
            CREATE VIEW raw_data AS 
            SELECT * FROM read_csv_auto('{file_path}', sample_size=-1)
        """)
        
        total_rows = conn.execute("SELECT COUNT(*) FROM raw_data").fetchone()[0]
        total_chunks = (total_rows // self.CHUNK_SIZE) + 1
        
        logger.info("chunk_info", total_rows=total_rows, total_chunks=total_chunks)
        
        # Step 2: Create output table
        conn.execute("""
            CREATE TABLE cleaned_data AS 
            SELECT * FROM raw_data WHERE 1=0
        """)
        
        issues_found = 0
        fixes_applied = 0
        
        # Step 3: Process chunks
        for chunk_idx in range(total_chunks):
            offset = chunk_idx * self.CHUNK_SIZE
            
            if progress_callback:
                progress_callback(
                    chunk_idx + 1, 
                    total_chunks,
                    f"Processing chunk {chunk_idx + 1}/{total_chunks}..."
                )
            
            logger.info("processing_chunk", 
                       chunk=chunk_idx + 1, 
                       total=total_chunks,
                       offset=offset)
            
            # Load chunk
            conn.execute(f"""
                CREATE TEMP TABLE chunk AS 
                SELECT * FROM raw_data 
                LIMIT {self.CHUNK_SIZE} OFFSET {offset}
            """)
            
            # Clean chunk (remove duplicates within chunk)
            chunk_rows = conn.execute("SELECT COUNT(*) FROM chunk").fetchone()[0]
            
            if chunk_rows == 0:
                break  # No more data
            
            # Remove duplicates in chunk
            duplicates = conn.execute("""
                SELECT COUNT(*) - COUNT(DISTINCT *) FROM chunk
            """).fetchone()[0]
            
            if duplicates > 0:
                issues_found += duplicates
                fixes_applied += duplicates
            
            # Insert cleaned chunk
            conn.execute("""
                INSERT INTO cleaned_data 
                SELECT DISTINCT * FROM chunk
            """)
            
            # Drop temp table
            conn.execute("DROP TABLE chunk")
        
        # Step 4: Final deduplication across all chunks
        if progress_callback:
            progress_callback(total_chunks, total_chunks, "Final deduplication...")
        
        final_count_before = conn.execute("SELECT COUNT(*) FROM cleaned_data").fetchone()[0]
        
        conn.execute("""
            CREATE TABLE final_data AS 
            SELECT DISTINCT * FROM cleaned_data
        """)
        
        final_count_after = conn.execute("SELECT COUNT(*) FROM final_data").fetchone()[0]
        cross_chunk_dupes = final_count_before - final_count_after
        
        if cross_chunk_dupes > 0:
            issues_found += cross_chunk_dupes
            fixes_applied += cross_chunk_dupes
            logger.info("cross_chunk_duplicates", count=cross_chunk_dupes)
        
        # Step 5: Handle nulls (optional, can be memory intensive)
        if not conservative:
            columns = conn.execute("PRAGMA table_info(final_data)").fetchall()
            for col_info in columns:
                col_name = col_info[1]
                
                nulls = conn.execute(f"""
                    SELECT COUNT(*) FROM final_data WHERE "{col_name}" IS NULL
                """).fetchone()[0]
                
                if nulls > 0:
                    issues_found += nulls
                    conn.execute(f"""
                        UPDATE final_data 
                        SET "{col_name}" = COALESCE("{col_name}", 'UNKNOWN')
                    """)
                    fixes_applied += nulls
        
        # Step 6: Convert to Polars DataFrame
        if progress_callback:
            progress_callback(total_chunks, total_chunks, "Converting to output format...")
        
        df = conn.execute("SELECT * FROM final_data").pl()
        
        quality_score = ((len(df) / total_rows) * 100) if total_rows > 0 else 100.0
        
        logger.info("chunked_processing_complete",
                   total_rows=total_rows,
                   cleaned_rows=len(df),
                   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
        )

