"""Data profiling for LLM decision-making.

Produces column-level statistics and a correlation matrix, with built-in
context‑window truncation so the output stays under a configurable budget
(default 16 KB).
"""
import json
from typing import Any, Optional

import polars as pl
import numpy as np


MAX_PROFILE_SIZE_KB = 16
"""Maximum serialised profile size before truncation kicks in."""

MAX_SAMPLE_ROWS = 5000
"""Maximum rows to collect for profiling."""

MAX_COLUMNS_FULL = 100
"""Beyond this column count the profile is always truncated."""

MAX_CORRELATION_COLS = 20
"""Maximum numeric columns to include in the correlation matrix."""

_ENTITY_KEYWORDS = frozenset({
    'ip', 'host', 'domain', 'user', 'client', 'src_ip', 'dst_ip',
    'source', 'destination', 'address', 'entity', 'name', 'id',
    'username', 'hostname', 'mac', 'email', 'phone', 'session',
})


def _is_float_dtype(dtype: pl.DataType) -> bool:
    return dtype in (pl.Float32, pl.Float64)


def _is_int_dtype(dtype: pl.DataType) -> bool:
    return dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                     pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64)


def _is_numeric_dtype(dtype: pl.DataType) -> bool:
    return _is_float_dtype(dtype) or _is_int_dtype(dtype)


def _is_string_dtype(dtype: pl.DataType) -> bool:
    return dtype in (pl.Utf8, pl.String, pl.Categorical)


def _is_temporal_dtype(dtype: pl.DataType) -> bool:
    return dtype in (pl.Date, pl.Datetime, pl.Time, pl.Duration)


# ── column statistics ───────────────────────────────────────────────────

def _compute_column_stats(series: pl.Series) -> dict:
    """Return a dict of statistics for a single column, based on its dtype.

    Numeric columns get min/max/mean/std/median/percentiles.
    String columns get unique count, missing count, and top-5 frequencies.
    Boolean columns get true/false/null counts.
    Temporal columns get min/max/null_count.
    """
    dtype = series.dtype
    stats: dict = {
        'dtype': str(dtype),
        'null_count': int(series.null_count()),
        'non_null_count': int(len(series) - series.null_count()),
    }

    if _is_numeric_dtype(dtype):
        non_null = series.drop_nulls()
        if len(non_null) > 0:
            stats.update({
                'min': _safe_float(non_null.min()),
                'max': _safe_float(non_null.max()),
                'mean': _safe_float(non_null.mean()),
                'std': _safe_float(non_null.std()),
                'median': _safe_float(non_null.median()),
                'q25': _safe_float(non_null.quantile(0.25)),
                'q75': _safe_float(non_null.quantile(0.75)),
                'unique_count': int(non_null.n_unique()),
            })
        else:
            stats.update({
                'min': None, 'max': None, 'mean': None, 'std': None,
                'median': None, 'q25': None, 'q75': None, 'unique_count': 0,
            })
        stats['type_category'] = 'numeric'

    elif _is_string_dtype(dtype):
        non_null = series.drop_nulls()
        stats['unique_count'] = int(non_null.n_unique()) if len(non_null) > 0 else 0
        stats['type_category'] = 'categorical'

        # Top-5 value counts
        if len(non_null) > 0:
            vc = non_null.value_counts().sort('count', descending=True).head(5)
            # Polars >=1.0 value_counts() uses original col name, not 'index'
            val_col = vc.columns[0]
            stats['top_values'] = [
                {'value': str(row[val_col]), 'count': int(row['count'])}
                for row in vc.iter_rows(named=True)
            ]
        else:
            stats['top_values'] = []

        # Detect entity column
        stats['is_entity_candidate'] = str(series.name).lower() in _ENTITY_KEYWORDS

    elif dtype == pl.Boolean:
        non_null = series.drop_nulls()
        if len(non_null) > 0:
            true_count = int(non_null.sum())
            stats.update({
                'true_count': true_count,
                'false_count': int(len(non_null)) - true_count,
            })
        else:
            stats.update({'true_count': 0, 'false_count': 0})
        stats['type_category'] = 'boolean'

    else:
        stats['type_category'] = 'other'
        stats['unique_count'] = int(series.drop_nulls().n_unique())

    return stats


def _safe_float(val: Any) -> Optional[float]:
    """Convert *val* to float or return None on failure."""
    if val is None:
        return None
    try:
        f = float(val)
        return None if np.isnan(f) or np.isinf(f) else round(f, 6)
    except (ValueError, TypeError):
        return None


# ── correlation matrix ──────────────────────────────────────────────────

def _compute_correlation_matrix(
    df: pl.DataFrame,
    top_n: int = MAX_CORRELATION_COLS,
) -> Optional[list[dict]]:
    """Pearson correlation for numeric columns.

    Returns a list of ``{column_i, column_j, correlation}`` dicts for the
    top *top_n* columns by mean absolute correlation, sorted descending by
    absolute value.  Returns *None* when fewer than 2 numeric columns exist.
    """
    numeric_cols = [
        col for col in df.columns
        if _is_numeric_dtype(df[col].dtype)
    ]
    if len(numeric_cols) < 2:
        return None

    # Limit to top_n numeric columns (pick by variance — highest first)
    if len(numeric_cols) > top_n:
        variances = []
        for col in numeric_cols:
            v = df[col].drop_nulls().std()
            variances.append((col, 0.0 if v is None else float(v)))
        variances.sort(key=lambda x: x[1], reverse=True)
        numeric_cols = [c for c, _ in variances[:top_n]]

    # Build correlation matrix
    try:
        import pandas as pd
        pdf = df.select(numeric_cols).to_pandas()
        corr = pdf.corr(method='pearson')
        pairs: list[dict] = []
        for i, col_i in enumerate(numeric_cols):
            for col_j in numeric_cols[i + 1:]:
                val = corr.loc[col_i, col_j]
                if np.isnan(val):
                    continue
                pairs.append({
                    'column_i': col_i,
                    'column_j': col_j,
                    'correlation': round(float(val), 4),
                })
    except ImportError:
        # Pandas not available — compute simple correlation matrix via numpy
        data = df.select(numeric_cols).to_numpy()
        data = np.nan_to_num(data, nan=0.0)
        corr_np = np.corrcoef(data.T)
        pairs = []
        for i in range(len(numeric_cols)):
            for j in range(i + 1, len(numeric_cols)):
                val = corr_np[i, j]
                if np.isnan(val):
                    continue
                pairs.append({
                    'column_i': numeric_cols[i],
                    'column_j': numeric_cols[j],
                    'correlation': round(float(val), 4),
                })

    # Sort by absolute correlation descending
    pairs.sort(key=lambda x: abs(x['correlation']), reverse=True)
    return pairs[:top_n * 2]  # cap output


# ── truncation ──────────────────────────────────────────────────────────

def _estimate_json_size(obj: Any) -> int:
    """Return the length of *obj* when serialised to JSON."""
    return len(json.dumps(obj, default=str, ensure_ascii=False))


def truncate_profile(
    profile: dict,
    max_kb: int = MAX_PROFILE_SIZE_KB,
) -> dict:
    """Remove low-value columns from *profile* until it fits in *max_kb*.

    Sets ``truncated=True`` and adds ``omitted_columns`` if truncation
    occurred.  Columns with the least variance (or highest null rate) are
    removed first.
    """
    max_bytes = max_kb * 1024
    if _estimate_json_size(profile) <= max_bytes:
        profile['truncated'] = False
        profile.setdefault('omitted_columns', [])
        return profile

    column_stats = profile.get('column_stats', {})
    if not column_stats:
        profile['truncated'] = True
        return profile

    # Score columns: low variance / high null rate → remove first
    scored: list[tuple[str, float]] = []
    for col_name, stats in column_stats.items():
        null_rate = stats.get('null_count', 0) / max(stats.get('non_null_count', 1), 1)
        # Numeric columns get a variance estimate
        std = stats.get('std')
        variance_score = float(std) if std is not None else 0.0
        # Remove low-variance OR high-null columns
        removal_score = -variance_score + null_rate * 100
        scored.append((col_name, removal_score))

    scored.sort(key=lambda x: x[1])  # most valuable first

    kept_columns = []
    omitted_columns = []
    for col_name, _ in scored:
        test_profile = dict(profile)
        test_profile['column_stats'] = {
            c: column_stats[c] for c in kept_columns + [col_name]
        }
        if _estimate_json_size(test_profile) <= max_bytes:
            kept_columns.append(col_name)
        else:
            omitted_columns.append(col_name)

    profile['column_stats'] = {
        c: column_stats[c] for c in kept_columns
    }
    profile['truncated'] = True
    profile['omitted_columns'] = omitted_columns
    profile['truncation_note'] = (
        f"{len(omitted_columns)} of {len(omitted_columns) + len(kept_columns)} "
        f"columns omitted to stay under {max_kb} KB budget"
    )
    return profile


# ── main profiling entry point ──────────────────────────────────────────

def profile_dataset(
    lf: pl.LazyFrame,
    sample_size: int = 1000,
    columns: Optional[list[str]] = None,
    max_kb: int = MAX_PROFILE_SIZE_KB,
    max_correlation_cols: int = MAX_CORRELATION_COLS,
) -> dict:
    """Compute column-level statistics and a correlation matrix.

    Args:
        lf: LazyFrame to profile.
        sample_size: Maximum rows to collect (sampled if larger).
        columns: Subset of columns to profile (``None`` = all).
        max_kb: Truncation budget in KB.
        max_correlation_cols: Max numeric columns for correlation.

    Returns:
        A dict with keys ``column_stats``, ``correlation_matrix``,
        ``truncated``, ``omitted_columns``, ``total_columns``,
        ``profiled_rows``, and ``sample_ratio``.
    """
    # Collect (possibly sampled) dataframe
    df = lf.collect(streaming=True)
    total_rows = len(df)
    sample_ratio = 1.0

    if total_rows > sample_size:
        df = df.sample(n=sample_size, seed=42)
        sample_ratio = sample_size / total_rows

    # Determine columns to profile
    if columns:
        profile_columns = [c for c in columns if c in df.columns]
    else:
        profile_columns = list(df.columns)

    if len(profile_columns) > MAX_COLUMNS_FULL:
        profile_columns = profile_columns[:MAX_COLUMNS_FULL]

    # Per-column stats
    column_stats: dict = {}
    for col_name in profile_columns:
        column_stats[col_name] = _compute_column_stats(df[col_name])

    # Correlation matrix
    correlation_matrix = _compute_correlation_matrix(df, top_n=max_correlation_cols)

    result = {
        'column_stats': column_stats,
        'correlation_matrix': correlation_matrix,
        'truncated': False,
        'omitted_columns': [],
        'total_columns': len(df.columns),
        'profiled_columns': len(profile_columns),
        'profiled_rows': len(df),
        'total_rows': total_rows,
        'sample_ratio': round(sample_ratio, 4),
    }

    return truncate_profile(result, max_kb=max_kb)