"""CSV loading with per-file schema validation, BOM detection, and multi-file merge.

Supports multiple encodings and delimiter-separated files.  The main entry
point is :func:`load_csv_directory`, which globs files, detects BOM per file,
validates schema consistency, and returns a merged :class:`pl.LazyFrame`.
"""
import glob
import logging
import os
import sqlite3
from pathlib import Path
from typing import Optional

import polars as pl
import yaml

from analysis.type_classifier import classify_schema, DataType

logger = logging.getLogger(__name__)


SUPPORTED_ENCODINGS = ['utf-8', 'utf-8-sig', 'utf-16le', 'utf-16be', 'latin-1']


# ── numeric coercion helper ──────────────────────────────────────────────

def _coerce_to_float(series: pl.Series) -> pl.Series:
    """Convert string series to Float64.

    Every entry that doesn't parse as a valid float becomes null.
    Handles blank strings, standalone '+'/'-', 'N/A', 'null', 'None', etc.
    """
    series = series.cast(pl.Utf8).str.strip_chars()
    artifacts = ('', '+', '-', '.', '..', 'N/A', 'n/a', 'NA', 'null', 'None', 'NULL', '?', '---', '--', '/')
    is_artifact = series.is_in(pl.Series(artifacts).implode())
    result = []
    for v, bad in zip(series.to_list(), is_artifact.to_list()):
        if v is None or bad:
            result.append(None)
        else:
            # Check for "+"/"-" prefix artifacts (e.g., "+abc", "-xyz")
            # where the value isn't a genuine numeric string.  Values like
            # "+5" or "-3.14" are still accepted as valid floats below.
            if len(v) > 1 and v[0] in '+-':
                try:
                    float(v[1:])  # validate suffix minus prefix
                except ValueError:
                    result.append(None)
                    continue
            try:
                result.append(float(v))
            except (ValueError, TypeError):
                result.append(None)
    return pl.Series(result, dtype=pl.Float64)


# ── BOM detection ───────────────────────────────────────────────────────

def detect_bom(path: str) -> str:
    """Read the first 4 bytes of *path* and return the detected encoding.

    Returns ``'utf-8'`` when no BOM is found.

    Note: Python 3 ``open()`` fully supports Unicode/Chinese paths on all
    platforms (no encoding dance needed).
    """
    with open(path, 'rb') as f:
        header = f.read(4)

    if header[:3] == b'\xef\xbb\xbf':
        return 'utf-8-sig'
    if header[:2] == b'\xff\xfe':
        return 'utf-16le'
    if header[:2] == b'\xfe\xff':
        return 'utf-16be'
    return 'utf-8'


# ── schema utilities ────────────────────────────────────────────────────

def collect_schema(
    path: str,
    encoding: str = 'utf-8',
    delimiter: str = ',',
    sample_rows: int = 10000,
) -> dict:
    """Return {column_name: dtype_string} for a single CSV file.

    Uses :func:`pl.scan_csv` with *sample_rows* rows for type inference.
    """
    lf = pl.scan_csv(
        str(path),  # str() ensures Unicode/Chinese paths resolve correctly
        encoding=encoding,
        separator=delimiter,
        infer_schema_length=sample_rows,
    )
    schema = lf.collect_schema()
    return {name: str(dtype) for name, dtype in zip(schema.names(), schema.dtypes())}


def _merge_schema_entries(
    schemas: list[dict],
) -> dict:
    """Merge multiple schema dicts, raising on type conflicts.

    Returns a single {name: dtype} dict.  If a column appears in multiple
    files with different dtypes the *last* non-null dtype wins (with a
    warning embedded in the result).
    """
    merged: dict[str, str] = {}
    conflicts: list[str] = []
    for s in schemas:
        for name, dtype in s.items():
            if name in merged and merged[name] != dtype:
                conflicts.append(name)
            merged[name] = dtype  # last wins
    return merged, conflicts


def validate_schemas(schemas: list[dict]) -> None:
    """Compare schemas across files and raise :class:`ValueError` on mismatch.

    Checks that every file has the same set of column names.  Type
    differences across files are **not** considered an error (Polars can
    cast), but differences in *column sets* are.
    """
    if not schemas:
        return

    col_sets = [set(s.keys()) for s in schemas]
    reference = col_sets[0]

    for i, cs in enumerate(col_sets[1:], start=1):
        only_in_ref = reference - cs
        only_in_current = cs - reference
        if only_in_ref or only_in_current:
            msg_parts = []
            if only_in_ref:
                msg_parts.append(f"File #{i} missing columns: {sorted(only_in_ref)}")
            if only_in_current:
                msg_parts.append(f"File #{i} has extra columns: {sorted(only_in_current)}")
            raise ValueError(f"Schema mismatch: {'; '.join(msg_parts)}")


# ── dtype inference from schema ─────────────────────────────────────────

def detect_schema_dtypes(
    lf: pl.LazyFrame,
    sample_rows: int = 10000,
) -> dict:
    """Infer Polars dtype strings from a LazyFrame's schema.

    The schema is read without collecting the full data (the scan scans
    a small sample).
    """
    schema = lf.collect_schema()
    return {name: str(dtype) for name, dtype in zip(schema.names(), schema.dtypes())}


# ── config loading ──────────────────────────────────────────────────────

def load_config(config_path: Optional[str]) -> dict:
    """Load a YAML config file; returns an empty dict if *config_path* is None."""
    if config_path is None:
        return {}
    path = Path(config_path)
    if not path.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")
    with open(path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f) or {}


# ── SQLite persistence ──────────────────────────────────────────────────

def _dtype_to_sqlite(dtype: pl.DataType) -> str:
    """Map a Polars dtype to a SQLite column type affinity."""
    if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                 pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
        return 'INTEGER'
    if dtype in (pl.Float32, pl.Float64):
        return 'REAL'
    if dtype == pl.Boolean:
        return 'INTEGER'
    return 'TEXT'  # String, Date, Datetime, etc. all map to TEXT


def save_to_db(run_id: int, lf: pl.LazyFrame, mode: str = 'replace') -> str:
    """Write a LazyFrame to a SQLite table ``_data_{run_id}``.

    Collects the LazyFrame (streaming if possible), creates or replaces the
    table, bulk-inserts rows in chunks, and returns the table name.

    Args:
        run_id: Analysis run ID (used as ``_data_{run_id}``).
        lf: LazyFrame to persist.
        mode: ``'replace'`` (default) → drop+create+insert.
              ``'append'`` → skip DDL, only INSERT (table must exist).

    This is a *persistence copy* — the original CSV files remain untouched.
    """
    from django.conf import settings

    table_name = f"_data_{run_id}"
    db_path = settings.DATABASES['default']['NAME']

    df = lf.collect(streaming=True)
    conn = sqlite3.connect(db_path)

    try:
        if mode == 'replace':
            # Drop existing table for retry/idempotency safety
            conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')

            # Build CREATE TABLE from Polars schema
            col_defs = [f'"{name}" {_dtype_to_sqlite(dtype)}'
                        for name, dtype in zip(df.columns, df.dtypes)]
            conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(col_defs)})')

        # Bulk insert in chunks
        col_names = [f'"{c}"' for c in df.columns]
        placeholders = ','.join(['?' for _ in df.columns])
        insert_sql = f'INSERT INTO "{table_name}" ({", ".join(col_names)}) VALUES ({placeholders})'

        total_rows = len(df)
        chunk_size = 10000
        for i in range(0, total_rows, chunk_size):
            chunk = df.slice(i, chunk_size)
            conn.executemany(insert_sql, chunk.rows())

        conn.commit()
    finally:
        conn.close()

    logger.info('[SAVE_TO_DB] table=%s rows=%d mode=%s', table_name, total_rows, mode)
    return table_name


_STR_TO_DTYPE = {
    'Float64': pl.Float64, 'Float32': pl.Float32,
    'Int64': pl.Int64, 'Int32': pl.Int32, 'Int16': pl.Int16, 'Int8': pl.Int8,
    'UInt64': pl.UInt64, 'UInt32': pl.UInt32, 'UInt16': pl.UInt16, 'UInt8': pl.UInt8,
    'Utf8': pl.Utf8, 'String': pl.Utf8, 'Categorical': pl.Categorical,
    'Boolean': pl.Boolean, 'Bool': pl.Boolean,
    'Date': pl.Date, 'Datetime': pl.Datetime, 'Time': pl.Time,
}


def _str_to_polars_dtype(s: str) -> pl.DataType:
    """Convert a Polars dtype string (e.g. ``'Float64'``, ``'Int32'``) to a ``pl.DataType``.

    Strips trailing metadata like ``Datetime(time_unit='us', time_zone=None)``.
    """
    base = s.split('(')[0].strip()
    return _STR_TO_DTYPE.get(base, pl.Utf8)


def load_from_db(
    table_name: str,
    schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
    """Load a SQLite table into a LazyFrame.

    Args:
        table_name: SQLite table name.
        schema_overrides: Optional dict of ``{column_name: dtype}`` to cast
            columns after loading.  Values can be ``pl.DataType`` instances
            (e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
            Use when SQLite ``TEXT`` columns need to be restored to
            ``Float64`` / ``Int64`` etc.

    Returns ``None`` if the table does not exist.
    """
    from django.conf import settings

    db_path = settings.DATABASES['default']['NAME']
    conn = sqlite3.connect(db_path)

    try:
        # Check table exists
        cursor = conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
            (table_name,),
        )
        if cursor.fetchone() is None:
            return None

        conn.row_factory = sqlite3.Row
        rows = conn.execute(f'SELECT * FROM "{table_name}"').fetchall()
        if not rows:
            return pl.DataFrame().lazy()
        try:
            df = pl.from_dicts([dict(r) for r in rows])
        except Exception:
            # Fallback: load all columns as Utf8 to handle mixed-type SQLite columns,
            # then apply schema_overrides to restore intended types.
            import io, csv as _csv_mod
            buf = io.StringIO()
            writer = _csv_mod.writer(buf)
            writer.writerow(rows[0].keys())
            for r in rows:
                writer.writerow(str(v) if v is not None else '' for v in r)
            buf.seek(0)
            df = pl.read_csv(buf, infer_schema_length=0)

        # Apply type overrides to restore original column types
        if schema_overrides:
            casts = {}
            for col, dtype in schema_overrides.items():
                if col not in df.columns:
                    continue
                dtype_obj = dtype if isinstance(dtype, pl.DataType) else _str_to_polars_dtype(dtype)
                if df[col].dtype != dtype_obj:
                    try:
                        casts[col] = pl.col(col).cast(dtype_obj)
                    except Exception:
                        # If a cast fails (e.g. "A" → Int64), keep the column as-is
                        pass
            if casts:
                df = df.with_columns(list(casts.values()))

        return df.lazy()
    finally:
        conn.close()


def drop_sqlite_table(table_name: str) -> None:
    """Drop a SQLite data table silently if it exists."""
    from django.conf import settings

    db_path = settings.DATABASES['default']['NAME']
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
        conn.commit()
    finally:
        conn.close()


# ── helpers ─────────────────────────────────────────────────────────────

def _file_count(paths: list[str]) -> int:
    return len(paths)


def _resolve_encoding(encoding: str, file_encoding: str) -> str:
    """BOM-detected encoding overrides user-supplied encoding.

    Polars >= 1.0 requires ``'utf8'`` (not ``'utf-8'``).
    """
    resolved = file_encoding if file_encoding != 'utf-8' else encoding
    return resolved.replace('-', '') if resolved in ('utf-8', 'utf-8-sig') else resolved


# ── main entry ──────────────────────────────────────────────────────────

def load_csv_directory(
    glob_pattern: str,
    encoding: str = 'utf-8',
    delimiter: str = ',',
    config_path: Optional[str] = None,
    schema_override: Optional[dict] = None,
    sample_rows: int = 10000,
    schema_strict: bool = False,
    recursive: bool = False,
    head: Optional[int] = None,
) -> tuple[pl.LazyFrame, dict, int, int, float]:
    """Load, validate, and merge CSV files matching *glob_pattern*.

    Args:
        glob_pattern: Glob pattern for CSV files (e.g. ``data/*.csv``).
        encoding: Fallback encoding when no BOM is detected.
        delimiter: Column delimiter character.
        config_path: Optional path to a YAML config (merged into metadata).
        schema_override: If given, force these {column: dtype} overrides.
        sample_rows: Rows to scan for dtype inference per file.
        schema_strict: When True, raise ValueError on column mismatch across
            files. When False (default), log warnings and merge with
            ``diagonal_relaxed`` (missing columns get null).
        recursive: When True, use ``glob.glob(..., recursive=True)`` to
            support ``**/*.csv`` patterns.
        head: Optional row limit for preview scenarios. When set, applies
            ``lf.head(head)`` before returning.

    Returns:
        A tuple of ``(lazyframe, schema_dict, total_row_count, file_count,
        memory_estimate_mb)``.

    Raises:
        FileNotFoundError: If no files match the glob pattern.
        ValueError: If schemas are inconsistent across files and
            *schema_strict* is True.
    """
    # Normalise path separators for cross-platform Unicode support
    glob_pattern = os.path.normpath(glob_pattern)

    # Resolve file list
    paths = sorted(glob.glob(glob_pattern, recursive=recursive))
    if not paths and not recursive:
        # Fallback: try recursive pattern
        paths = sorted(glob.glob(glob_pattern, recursive=True))
    if not paths:
        raise FileNotFoundError(f"No files match glob pattern: {glob_pattern}")

    # Load config (user-provided or fall back to project default)
    config = load_config(config_path)
    if not config:
        default_path = Path(__file__).parent.parent / 'config' / 'config.yaml'
        if default_path.exists():
            config = load_config(str(default_path))

    # Collect per-file schema
    schemas: list[dict] = []
    lazyframes: list[pl.LazyFrame] = []
    for path in paths:
        file_encoding = detect_bom(path)
        actual_encoding = _resolve_encoding(encoding, file_encoding)

        lf = pl.scan_csv(
            str(path),  # str() ensures Unicode/Chinese paths resolve correctly
            encoding=actual_encoding,
            separator=delimiter,
            infer_schema_length=max(sample_rows, 1000),
        )
        schema = detect_schema_dtypes(lf, sample_rows=sample_rows)

        # Apply schema override for this file
        if schema_override:
            for col, dtype in schema_override.items():
                if col in schema:
                    schema[col] = dtype

        schemas.append(schema)
        lazyframes.append(lf)

    if schema_strict:
        # Raise on column mismatch
        validate_schemas(schemas)
    else:
        # Lenient mode: log warnings for column mismatches
        col_sets = [set(s.keys()) for s in schemas]
        reference = col_sets[0]
        for i, cs in enumerate(col_sets[1:], start=1):
            only_in_ref = reference - cs
            only_in_current = cs - reference
            if only_in_ref:
                logger.warning(
                    "File #%d missing columns: %s (will be null)",
                    i + 1, sorted(only_in_ref),
                )
            if only_in_current:
                logger.warning(
                    "File #%d has extra columns: %s (will be null in other files)",
                    i + 1, sorted(only_in_current),
                )

    # Merge schemas (unify column ordering — first file wins order)
    merged_schema, _ = _merge_schema_entries(schemas)

    # Concatenate all LazyFrames (union vertically)
    merged_lf: pl.LazyFrame = pl.concat(lazyframes, how='diagonal_relaxed')

    # Log schema and size estimate
    total_row_count_est = len(paths) * (lazyframes[0].collect(streaming=True).height if lazyframes else 0)
    logger.info(f'[LOAD] files={len(paths)} total_rows_est={total_row_count_est}')
    for col_name, dtype in sorted(merged_schema.items()):
        logger.info(f'[LOAD_SCHEMA] col={col_name} dtype={dtype}')

    # Apply schema override on merged frame (cast columns)
    if schema_override:
        casts = {}
        for col, dtype_str in schema_override.items():
            if col in merged_schema:
                try:
                    polars_dtype = getattr(pl, dtype_str.split('(')[0], None)
                    if polars_dtype:
                        casts[col] = pl.col(col).cast(polars_dtype)
                except Exception:
                    pass  # skip uncastable overrides silently
        if casts:
            merged_lf = merged_lf.with_columns(list(casts.values()))

    # ── Cleaning: drop outdated columns ──────────────────────────────────
    geoip_cfg = config.get('data', {}).get('geoip', {})
    drop_cols = geoip_cfg.get(
        'drop_columns',
        [],  # Empty by default — lat/lon columns are now cleaned, not dropped
    )
    existing_drop = [c for c in drop_cols if c in merged_schema]
    if existing_drop:
        merged_lf = merged_lf.drop(existing_drop)
        merged_schema = {k: v for k, v in merged_schema.items() if k not in existing_drop}
        logger.info('[CLEAN] dropped columns: %s', existing_drop)

    # ── Cleaning: BOOL_ENUM normalisation ────────────────────────────────
    try:
        type_map = classify_schema(
            merged_lf.head(1000),
            config_overrides=config.get('columns', {}),
        )
    except Exception:
        type_map = {}

    bool_enum_cols = [
        c for c, t in type_map.items()
        if t == DataType.BOOL_ENUM and c in merged_schema
    ]
    if bool_enum_cols:
        # Only clean string-typed BOOL_ENUM columns; native Boolean is already clean
        schema_dtypes = dict(zip(
            merged_lf.collect_schema().names(),
            merged_lf.collect_schema().dtypes(),
        ))
        bool_exprs = []
        cleaned_names: list[str] = []
        for c in bool_enum_cols:
            if schema_dtypes.get(c) not in (pl.Utf8, pl.String):
                continue  # native Boolean / numeric — skip
            c_str = pl.col(c).cast(pl.Utf8).str.strip_chars()
            cleaned = (
                pl.when(c_str == "")
                .then(pl.lit(None, pl.Utf8))
                .when(c_str.str.to_lowercase() == "+")
                .then(pl.lit("True"))
                .when(c_str.str.to_lowercase() == "-")
                .then(pl.lit("False"))
                .otherwise(c_str)
            )
            bool_exprs.append(cleaned.alias(c))
            cleaned_names.append(c)
        if bool_exprs:
            merged_lf = merged_lf.with_columns(bool_exprs)
            for c in cleaned_names:
                merged_schema[c] = 'Utf8'
            logger.info('[CLEAN] BOOL_ENUM columns normalised: %s', cleaned_names)

    # ── Cleaning: LAT_LON normalisation ────────────────────────────────
    latlon_cols = [
        c for c, t in type_map.items()
        if t == DataType.LAT_LON and c in merged_schema
    ]
    if latlon_cols:
        schema_dtypes_ll = dict(zip(
            merged_lf.collect_schema().names(),
            merged_lf.collect_schema().dtypes(),
        ))
        latlon_exprs = []
        for c in latlon_cols:
            dtype = schema_dtypes_ll.get(c)
            if dtype in (pl.Utf8, pl.String):
                latlon_exprs.append(
                    pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
                )
                merged_schema[c] = 'Float64'
        if latlon_exprs:
            merged_lf = merged_lf.with_columns(latlon_exprs)
            logger.info('[CLEAN] LAT_LON columns normalised to Float64: %s', latlon_cols)

    # ── Cleaning: Generic numeric coercion ────────────────────────────
    # For any remaining string column (not IPv4/URL/HEX/ENUM/etc.),
    # try to coerce to Float64. Whatever can't parse becomes null.
    # This prevents "mean on str column" crashes for ANY column.
    _SKIP_NUM_COERCE = (DataType.IPv4, DataType.URL, DataType.HEX,
                         DataType.ENUM, DataType.LAT_LON)
    # Also protect columns with these name patterns (MAC addresses, etc.)
    _SKIP_NAME_PATTERNS = ('mac', 'mac_', '_mac', 'addr', 'address')
    schema_dtypes_gn = dict(zip(
        merged_lf.collect_schema().names(),
        merged_lf.collect_schema().dtypes(),
    ))
    generic_numeric_cols = []
    for c, t in type_map.items():
        if t in _SKIP_NUM_COERCE:
            continue
        if c not in merged_schema:
            continue
        dtype = schema_dtypes_gn.get(c)
        if dtype not in (pl.Utf8, pl.String):
            continue  # already numeric
        # Skip columns with protected name patterns
        cl = c.lower().replace('-', '_').replace(' ', '_')
        if any(p in cl for p in _SKIP_NAME_PATTERNS):
            continue
        generic_numeric_cols.append(c)

    if generic_numeric_cols:
        num_exprs = []
        for c in generic_numeric_cols:
            num_exprs.append(
                pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
            )
        if num_exprs:
            merged_lf = merged_lf.with_columns(num_exprs)
            for c in generic_numeric_cols:
                merged_schema[c] = 'Float64'
            logger.info('[CLEAN] Generic numeric coercion applied: %s', generic_numeric_cols)

    # ── Cleaning: HEX preprocessing (hex_pairs_count) ────────────────────
    hex_cols = [
        c for c, t in type_map.items()
        if t == DataType.HEX and c in merged_schema
    ]
    if hex_cols:
        hex_exprs = []
        for c in hex_cols:
            hex_exprs.append(
                pl.col(c).str.split(" ").list.len().alias(f"{c}_hex_pairs_count")
            )
        merged_lf = merged_lf.with_columns(hex_exprs)
        for c in hex_cols:
            merged_schema[f"{c}_hex_pairs_count"] = 'UInt32'
        logger.info('[CLEAN] HEX columns added hex_pairs_count: %s', hex_cols)

    # Row count estimate (fast path: use first file row count × files)
    # For accurate count we'd need to collect, but that defeats lazy.
    # We estimate from the first file's scanned schema info.
    first_rows = lazyframes[0].collect(streaming=True).height if lazyframes else 0
    total_row_count = first_rows * len(paths)

    # Memory estimate
    bytes_per_row = 0
    for col_name, dtype_str in merged_schema.items():
        numeric_part = dtype_str.split('(')[0].lower()
        size_map = {
            'int8': 1, 'int16': 2, 'int32': 4, 'int64': 8,
            'uint8': 1, 'uint16': 2, 'uint32': 4, 'uint64': 8,
            'float32': 4, 'float64': 8,
            'boolean': 1, 'bool': 1,
            'utf8': 50, 'string': 50, 'str': 50,
            'categorical': 8, 'cat': 8,
            'date': 4, 'datetime': 8, 'time': 8, 'duration': 8,
        }
        bytes_per_row += size_map.get(numeric_part, 8)
    memory_mb = (bytes_per_row * total_row_count) / (1024 * 1024)

    # Apply head limit for preview (keeps LazyFrame lazy — rows are limited on collect)
    if head is not None:
        merged_lf = merged_lf.head(head)
        total_row_count = min(total_row_count, head)

    return merged_lf, merged_schema, total_row_count, len(paths), memory_mb