"""High-performance JSON-to-database pipeline using Polars vectorized transformations."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/09_utils_polars_mapper.ipynb.

# %% ../nbs/09_utils_polars_mapper.ipynb #386a8a41
from __future__ import annotations
import polars as pl
from sqlalchemy import create_engine, text
from typing import Dict, List, Optional
import uuid
import logging
from nbdev.showdoc import show_doc


logger = logging.getLogger(__name__)

# %% auto #0
__all__ = ['logger', 'map_and_upsert', 'apply_schema']

# %% ../nbs/09_utils_polars_mapper.ipynb #eb43d10e
def map_and_upsert(
    df: pl.DataFrame, # The raw Polars DataFrame from JSON
    table_name: str, # Target database table name
    key_col: str, # Primary key column for conflict resolution
    db_uri: str, # SQLAlchemy connection string (e.g., 'sqlite:///db.db' or 'postgresql://...')
    column_map: dict = None, # Optional rename map {json_key: db_col}
    unnest_cols: list[str] = None, # List of Struct columns to flatten
    type_map: dict = None # Optional type casting map {col_name: pl.DataType}
) -> int:
    """
    Map JSON data to database columns and upsert using staging table pattern.
    
    **Returns:** Number of rows affected by the upsert operation
    
    **Type Casting:**
    When columns have `None` values, Polars may infer incorrect types (e.g., String
    instead of Int64). This causes PostgreSQL type mismatch errors. Use `type_map`
    to explicitly cast columns before writing to the database.
    """
    # Step 1: Rename columns if mapping provided
    if column_map:
        df = df.rename(column_map)
        logger.info(f"Renamed columns: {column_map}")
    
    # Step 2: Flatten nested columns if specified
    if unnest_cols:
        for col in unnest_cols:
            if col in df.columns:
                df = df.unnest(col)
                logger.info(f"Unnested column: {col}")
    
    # Step 3: Apply type casting if specified
    if type_map:
        cast_count = 0
        for col_name, dtype in type_map.items():
            if col_name not in df.columns:
                logger.warning(f"Column '{col_name}' in type_map not found in DataFrame, skipping")
                continue
            try:
                df = df.with_columns(pl.col(col_name).cast(dtype, strict=False))
                cast_count += 1
                logger.debug(f"Cast column '{col_name}' to {dtype}")
            except Exception as e:
                logger.warning(f"Failed to cast column '{col_name}' to {dtype}: {e}")
        logger.info(f"Applied type casting to {cast_count}/{len(type_map)} columns")
    
    # Step 4: Select only columns that exist in target table (drop extras)
    # This prevents errors from extra JSON fields
    engine = create_engine(db_uri)
    
    # Get target table columns
    with engine.connect() as conn:
        result = conn.execute(text(f"SELECT * FROM {table_name} LIMIT 0"))
        target_columns = list(result.keys())
    
    # Filter DataFrame to only target columns
    available_cols = [col for col in target_columns if col in df.columns]
    df = df.select(available_cols)
    logger.info(f"Selected columns for {table_name}: {available_cols}")
    
    # Step 5: Generate unique staging table name
    staging_table = f"staging_{uuid.uuid4().hex[:8]}"
    rows_affected = 0
    
    try:
        # Step 6: Write to staging table (fast bulk insert)
        df.write_database(
            table_name=staging_table,
            connection=db_uri,
            if_table_exists='replace'
        )
        logger.info(f"Wrote {len(df)} rows to staging table {staging_table}")
        
        # Step 7: Determine database type for dialect-specific SQL
        is_sqlite = 'sqlite' in db_uri.lower()
        
        # Step 8: Execute upsert from staging to target
        with engine.connect() as conn:
            if is_sqlite:
                # SQLite: INSERT OR REPLACE
                cols_str = ', '.join(available_cols)
                upsert_sql = f"""
                    INSERT OR REPLACE INTO {table_name} ({cols_str})
                    SELECT {cols_str} FROM {staging_table}
                """
            else:
                # PostgreSQL: INSERT ... ON CONFLICT DO UPDATE
                cols_str = ', '.join(available_cols)
                update_cols = [col for col in available_cols if col != key_col]
                update_set = ', '.join([f"{col} = EXCLUDED.{col}" for col in update_cols])
                
                upsert_sql = f"""
                    INSERT INTO {table_name} ({cols_str})
                    SELECT {cols_str} FROM {staging_table}
                    ON CONFLICT ({key_col}) DO UPDATE SET {update_set}
                """
            
            result = conn.execute(text(upsert_sql))
            rows_affected = result.rowcount if result.rowcount >= 0 else len(df)
            conn.commit()
            logger.info(f"Upserted {rows_affected} rows into {table_name}")
    
    finally:
        # Step 9: Cleanup - drop staging table
        with engine.connect() as conn:
            conn.execute(text(f"DROP TABLE IF EXISTS {staging_table}"))
            conn.commit()
            logger.debug(f"Dropped staging table {staging_table}")
    
    return rows_affected

# %% ../nbs/09_utils_polars_mapper.ipynb #707f8f07
def apply_schema(
    df: pl.DataFrame, # Input DataFrame
    type_map: dict # Column name -> Polars dtype (e.g., {'created_at': pl.Date, 'is_active': pl.Boolean})
) -> pl.DataFrame:
    """Apply explicit type conversions to DataFrame columns."""
    conversions = []
    
    for col_name, dtype in type_map.items():
        if col_name not in df.columns:
            logger.warning(f"Column {col_name} not found in DataFrame, skipping")
            continue
        
        # Handle different type conversions
        if dtype == pl.Date:
            conversions.append(pl.col(col_name).str.strptime(pl.Date, "%Y-%m-%d").alias(col_name))
        elif dtype == pl.Datetime:
            conversions.append(pl.col(col_name).str.strptime(pl.Datetime).alias(col_name))
        elif dtype == pl.Boolean:
            # Handle "true"/"false" strings
            conversions.append(
                pl.col(col_name).str.to_lowercase().eq("true").alias(col_name)
            )
        else:
            # Cast to specified type (works for numeric types)
            conversions.append(pl.col(col_name).cast(dtype).alias(col_name))
    
    if conversions:
        df = df.with_columns(conversions)
        logger.info(f"Applied schema conversions to {len(type_map)} columns")
    
    return df
