"""Entity profile building + adaptive scoring handlers."""
import traceback
from analysis.constants import MAX_DISTRIBUTION_SAMPLE, RANDOM_SEED
import logging
from typing import Optional

import numpy as np
import polars as pl

from ..session_store import SessionStore, _generate_id
from ..profile_util import profile
from ._helpers import _resolve_dataset, _count_rows, NUMERIC_DTYPES
logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════════════
#  Constants
# ═══════════════════════════════════════════════════════════════════════

_ENTITY_KEYWORDS = [
    'src_ip', 'dst_ip', 'source_ip', 'destination_ip', 'src_addr', 'dst_addr',
    'sip', 'dip', 'ip_src', 'ip_dst', 'saddr', 'daddr',
    'sni', 'host', 'domain', 'uri', 'url', 'fqdn', 'server_name',
    'mac', 'mac_addr', 'src_mac', 'dst_mac',
    'user_agent', 'ua', 'asn', 'organization', 'org',
]

_NUMERIC_TYPES_FOR_AGG = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                           pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                           pl.Float32, pl.Float64)

_ANOMALY_FEATURE_COLS = [
    'non_std_port_rate', 'modern_tls_rate', 'sni_missing_ratio',
    'recoverable_ratio', 'flow_count', 'total_bytes_sent',
    'total_packets', 'unique_dst_ips', 'unique_dst_ports',
    'unique_protocols', 'has_sni', 'avg_sni_length',
    'countries_visited',
]


# ═══════════════════════════════════════════════════════════════════════
#  _handle_build_entity_profiles
# ═══════════════════════════════════════════════════════════════════════

async def _handle_build_entity_profiles(
    dataset_id: str,
    entity_column: Optional[str] = None,
    entity_columns: Optional[list[str]] = None,
    auto_detect: bool = True,
) -> dict:
    """Group raw flow data by entity column(s) and compute per-entity aggregate features."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    schema = entry.get('schema', {})
    col_names = list(schema.keys())

    # ── Resolve entity columns ──────────────────────────────────────────
    if entity_column and entity_column in col_names:
        group_cols = [entity_column]
    elif entity_columns:
        group_cols = [c for c in entity_columns if c in col_names]
        if not group_cols:
            return {'error': f'None of entity_columns {entity_columns} exist in dataset. Available: {col_names[:20]}',
                    'truncated': False}
    elif auto_detect:
        matched = []
        for kw in _ENTITY_KEYWORDS:
            found = [c for c in col_names if kw.lower() in c.lower()]
            matched.extend(found)
        seen = set()
        group_cols = [c for c in matched if not (c in seen or seen.add(c))]
        if not group_cols:
            for c in col_names:
                if c.startswith('_'):
                    continue
                group_cols = [c]
                break
        group_cols = group_cols[:3]
    else:
        return {'error': 'No entity column specified and auto_detect is disabled',
                'truncated': False}

    if not group_cols:
        return {'error': 'Could not determine entity column(s) for grouping',
                'truncated': False}

    # ── Build aggregation expressions ──────────────────────────────────
    agg_exprs: list[pl.Expr] = [
        pl.len().alias('flow_count'),
        pl.col(group_cols[0]).n_unique().alias(f'unique_{group_cols[0]}'),
    ]

    numeric_cols = [c for c in col_names if c not in group_cols
                    and schema[c].startswith(('Int', 'UInt', 'Float'))]
    for c in numeric_cols[:20]:
        agg_exprs.append(pl.sum(c).alias(f'total_{c}'))
        agg_exprs.append(pl.mean(c).alias(f'avg_{c}'))

    str_cols = [c for c in col_names if c not in group_cols
                and schema[c].startswith(('Utf8', 'String', 'Cat'))]
    for c in str_cols[:10]:
        agg_exprs.append(pl.col(c).n_unique().alias(f'unique_{c}'))

    for kw in ['dst_ip', 'dip', 'dst_addr', 'daddr', 'ip_dst', 'dest_ip',
               'dst_port', 'dport', 'destination_port', 'dest_port',
               'port', 'protocol', 'proto', 'sni', 'server_name',
               'cipher_suite', 'cipher', 'tls_version', 'version', '0ver',
               'country', 'scnt', 'dcnt', 'asn', 'organization']:
        found = [c for c in col_names if c not in group_cols and kw.lower() in c.lower()]
        for c in found:
            alias_name = f'unique_{c}'
            if alias_name not in [str(a.meta.output_name()) if hasattr(a, 'meta') else '' for a in agg_exprs]:
                agg_exprs.append(pl.col(c).n_unique().alias(alias_name))

    # ── Execute grouping ───────────────────────────────────────────────
    row_count = _count_rows(lf)
    ENTITY_CHUNK_THRESHOLD = 1_000_000   # 1M rows
    ENTITY_CHUNK_SIZE = 100_000          # 100K rows per chunk

    if row_count >= ENTITY_CHUNK_THRESHOLD:
        # Chunked path: process in 100K batches → group → merge across chunks
        all_chunk_dfs = []
        for offset in range(0, row_count, ENTITY_CHUNK_SIZE):
            chunk_df = lf.slice(offset, ENTITY_CHUNK_SIZE).collect(streaming=True)
            chunk_grouped = chunk_df.group_by(group_cols).agg(agg_exprs)
            all_chunk_dfs.append(chunk_grouped)
        merged_df = pl.concat(all_chunk_dfs)

        # Re-aggregate across chunks (same entity may appear in multiple chunks)
        re_agg_exprs: list[pl.Expr] = []
        for col_name in merged_df.columns:
            if col_name in group_cols:
                continue
            if col_name == 'flow_count':
                re_agg_exprs.append(pl.sum('flow_count').alias('flow_count'))
            elif col_name.startswith('total_'):
                re_agg_exprs.append(pl.sum(col_name).alias(col_name))
            elif col_name.startswith('avg_'):
                # Weighted average: sum(total) / sum(flow_count)
                base = col_name[4:]
                total_col = f'total_{base}'
                if total_col in merged_df.columns:
                    re_agg_exprs.append(
                        (pl.sum(total_col) / pl.sum('flow_count')).alias(col_name))
            elif col_name.startswith('unique_'):
                # Approximate: max of per-chunk distinct counts
                re_agg_exprs.append(pl.max(col_name).alias(col_name))

        entity_lf = merged_df.lazy().group_by(group_cols).agg(re_agg_exprs)
        logger.info('[ENTITY] chunked aggregation: %d rows → %d chunks → %d merged rows',
                    row_count, len(all_chunk_dfs), len(merged_df))
    else:
        entity_lf = lf.group_by(group_cols).agg(agg_exprs)

    n_entities = _count_rows(entity_lf)

    # ── Store result ───────────────────────────────────────────────────
    new_schema = {name: str(dtype) for name, dtype in
                  zip(entity_lf.collect_schema().names(), entity_lf.collect_schema().dtypes())}
    new_id = f'entity_{_generate_id("ds")}'
    store = SessionStore()
    store.store_dataset(
        dataset_id=new_id,
        lazyframe=entity_lf,
        schema=new_schema,
        metadata={
            'row_count': n_entities,
            'entity_columns': group_cols,
            'parent_dataset_id': dataset_id,
            'n_agg_features': len(agg_exprs),
        },
        parent_id=dataset_id,
    )

    return {
        'dataset_id': new_id,
        'entity_columns': group_cols,
        'n_entities': n_entities,
        'n_features': len(agg_exprs),
        'aggregated_columns': list(new_schema.keys())[:20],
    }


# ═══════════════════════════════════════════════════════════════════════
#  Adaptive scoring helpers
# ═══════════════════════════════════════════════════════════════════════

@profile
def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame:
    """Add anomaly scores using data-driven weights (SVD + IQR + percentile)."""
    _agg_cols = lf.collect_schema().names()
    avail = [c for c in _ANOMALY_FEATURE_COLS if c in _agg_cols]
    if not avail:
        return lf

    try:
        n_rows = _count_rows(lf)
        sample_lf = lf if n_rows <= MAX_DISTRIBUTION_SAMPLE else lf.sample(n=MAX_DISTRIBUTION_SAMPLE, seed=RANDOM_SEED)
        df_sample = sample_lf.select(avail).collect(streaming=True)
        mat = df_sample.to_numpy()
        mat = np.nan_to_num(mat, nan=0.0)

        from sklearn.decomposition import TruncatedSVD
        n_comp = min(1, min(mat.shape) - 1)
        if n_comp < 1 or mat.shape[1] < 2:
            return lf
        svd = TruncatedSVD(n_components=1, random_state=RANDOM_SEED).fit(mat)
        raw_weights = np.abs(svd.components_[0])
        weights = dict(zip(avail, raw_weights / raw_weights.sum()))

        expr_parts = []
        for col, w in weights.items():
            if col == 'modern_tls_rate':
                expr_parts.append((1.0 - pl.col(col).fill_null(0.5)) * w)
            else:
                expr_parts.append(pl.col(col).fill_null(0.0) * w)
        if expr_parts:
            proxy_expr = expr_parts[0]
            for p in expr_parts[1:]:
                proxy_expr = proxy_expr + p
        else:
            return lf

        lf = lf.with_columns(proxy_expr.clip(0, 1).alias('_raw_score'))

        sample_scores = df_sample.select(
            pl.sum_horizontal([
                (1.0 - pl.col('modern_tls_rate').fill_null(0.5)) * weights.get('modern_tls_rate', 0)
                if c == 'modern_tls_rate'
                else pl.col(c).fill_null(0.0) * weights.get(c, 0)
                for c in avail
            ]).clip(0, 1)
        ).to_series().to_numpy()

        q25, q50, q75 = np.percentile(sample_scores, [25, 50, 75])
        iqr = q75 - q25
        upper_fence = float(q75 + 1.5 * iqr) if iqr > 0 else 0.9
        extreme_fence = float(q75 + 3.0 * iqr) if iqr > 0 else 0.99

        lf = lf.with_columns([
            pl.col('_raw_score').alias('proxy_score'),
            pl.when(pl.col('_raw_score') >= extreme_fence)
            .then(pl.lit('critical'))
            .when(pl.col('_raw_score') >= upper_fence)
            .then(pl.lit('suspicious'))
            .when(pl.col('_raw_score') >= q50)
            .then(pl.lit('watch'))
            .otherwise(pl.lit('normal'))
            .alias('risk_level'),
            (pl.col('_raw_score') * 0.6 + 0.4 * (pl.col('_raw_score') > q75).cast(pl.Float64))
            .alias('threat_score'),
        ]).drop('_raw_score')

        valid_counts = {c: int((mat[:, i] != 0).sum()) for i, c in enumerate(avail)}
        logger = logging.getLogger(__name__)
        logger.info(
            '[SCORE] adaptive weights from SVD: %s | thresholds: p50=%.3f upper=%.3f extreme=%.3f | sample=%d/%d',
            {c: f'{w:.3f}' for c, w in sorted(weights.items(), key=lambda x: -x[1])},
            float(q50), upper_fence, extreme_fence,
            min(n_rows, MAX_DISTRIBUTION_SAMPLE), n_rows,
        )

    except Exception as exc:
        logging.getLogger(__name__).warning('[SCORE] adaptive scoring skipped: %s', exc)

    return lf


# ═══════════════════════════════════════════════════════════════════════
#  _handle_compute_scores
# ═══════════════════════════════════════════════════════════════════════

async def _handle_compute_scores(dataset_id: str) -> dict:
    """Compute adaptive proxy/anomaly/threat scores on entity data."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf = entry['lazyframe']
    lf = _add_adaptive_scores(lf, [])
    try:
        updated_schema = lf.collect_schema()
        new_schema = {name: str(dtype) for name, dtype in
                      zip(updated_schema.names(), updated_schema.dtypes())}
    except Exception:
        logger.error("_handle_compute_scores failed: {}".format(traceback.format_exc()))
        new_schema = entry.get('schema', {})
    store = SessionStore()
    store.store_dataset(dataset_id, lf, schema=new_schema,
                        metadata=entry.get('metadata', {}))
    return {'status': 'scored', 'dataset_id': dataset_id}
