"""Tool definitions and handler functions for the TLS Analyzer MCP server.

Each tool is defined by a :class:`Tool <mcp.types.Tool>` metadata object and
an async handler function.  Handlers interact with the :class:`SessionStore`
singleton and the Django ORM.

All 12 tools:
   1. load_data          7. filter_and_cluster
   2. profile_data       8. build_entity_profiles
   3. filter_data        9. export_results
   4. preprocess_data   10. list_datasets
   5. run_clustering    11. drop_dataset
   6. evaluate_clustering   12. clone_dataset
"""
import json
import os
from pathlib import Path
from typing import Any, Optional

import numpy as np
import polars as pl
from mcp.types import Tool

from .data_loader import load_csv_directory
from .data_profiler import profile_dataset as _profile_dataset
from .session_store import SessionStore, _generate_id


# ═══════════════════════════════════════════════════════════════════════
#  Tool metadata
# ═══════════════════════════════════════════════════════════════════════

def _json_size(obj: Any) -> int:
    return len(json.dumps(obj, default=str, ensure_ascii=False))


_MAX_RESPONSE_KB = 64
_MAX_RESPONSE_BYTES = _MAX_RESPONSE_KB * 1024


def _truncate_response(result: dict, max_kb: int = _MAX_RESPONSE_KB) -> dict:
    """Truncate large lists/keys in *result* to stay under *max_kb*."""
    if _json_size(result) <= max_kb * 1024:
        result.setdefault('truncated', False)
        return result

    # Drop verbose keys first
    for costly_key in ('column_stats', 'features', 'correlation_matrix',
                        'labels', 'cluster_labels'):
        if costly_key in result and isinstance(result[costly_key], dict):
            result[costly_key] = {'__truncated__': True,
                                  'key_count': len(result[costly_key])}
        elif costly_key in result and isinstance(result[costly_key], list):
            lst = result[costly_key]
            if len(lst) > 10:
                result[costly_key] = lst[:10] + [{'__truncated__': True,
                                                   'omitted': len(lst) - 10}]
    result['truncated'] = True
    return result


def get_tools_meta() -> list[Tool]:
    """Return the list of Tool metadata objects for MCP server registration.

    Call this from :meth:`TLSAnalyzerMCPServer._setup_tools`.
    """
    return [
        Tool(
            name="load_data",
            description=(
                "加载匹配glob模式的CSV文件。自动检测每个文件的BOM，验证schema一致性，"
                "纵向合并，并将结果作为惰性数据集存储在会话存储中。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "csv_glob": {
                        "type": "string",
                        "description": "Glob pattern for CSV files (e.g. 'data/*.csv')",
                    },
                    "config_path": {
                        "type": "string",
                        "description": "Optional path to a YAML config file",
                    },
                    "encoding": {
                        "type": "string",
                        "description": "Fallback encoding when no BOM is found",
                        "default": "utf-8",
                    },
                    "delimiter": {
                        "type": "string",
                        "description": "Column delimiter",
                        "default": ",",
                    },
                    "schema_strict": {
                        "type": "boolean",
                        "description": "When True, raise on column mismatch; when False, merge leniently with nulls",
                        "default": False,
                    },
                    "recursive": {
                        "type": "boolean",
                        "description": "When True, use recursive glob for **/*.csv patterns",
                        "default": False,
                    },
                },
                "required": ["csv_glob"],
            },
        ),
        Tool(
            name="profile_data",
            description=(
                "计算数据集的每列统计信息和数值相关矩阵。结果自动截断至约16 KB以"
                "节省LLM上下文窗口空间。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Dataset identifier from a previous load_data call",
                    },
                    "sample_size": {
                        "type": "integer",
                        "description": "Maximum rows to profile (sampled if larger)",
                        "default": 1000,
                    },
                    "columns": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Columns to profile (default: all)",
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="filter_data",
            description=(
                "对数据集应用列级过滤器，将结果存储为新的派生数据集。支持eq/neq/gt/gte/"
                "lt/lte/contains/not_contains/in/not_in/is_null/is_not_null运算符。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Source dataset identifier",
                    },
                    "filters": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "column": {"type": "string"},
                                "op": {
                                    "type": "string",
                                    "enum": [
                                        "eq", "neq", "gt", "gte", "lt", "lte",
                                        "contains", "not_contains",
                                        "in", "not_in",
                                        "is_null", "is_not_null",
                                    ],
                                },
                                "value": {"description": "Operand value"},
                            },
                            "required": ["column", "op"],
                        },
                        "description": "List of filter expressions",
                    },
                    "logic": {
                        "type": "string",
                        "enum": ["and", "or"],
                        "description": "Combine filters with AND or OR logic",
                        "default": "and",
                    },
                },
                "required": ["dataset_id", "filters"],
            },
        ),
        Tool(
            name="preprocess_data",
            description=(
                "预处理选定列：标准化（StandardScaler）、独热编码、缺失值填充和列删除。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Source dataset identifier",
                    },
                    "columns": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Columns to preprocess",
                    },
                    "config": {
                        "type": "object",
                        "properties": {
                            "standardize": {
                                "type": "boolean",
                                "description": "Apply StandardScaler to numeric columns",
                                "default": False,
                            },
                            "onehot": {
                                "type": "boolean",
                                "description": "One-hot encode categorical columns",
                                "default": False,
                            },
                            "fillna": {
                                "type": "string",
                                "enum": ["mean", "median", "zero", "drop"],
                                "description": "Missing value strategy",
                            },
                            "drop": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "Columns to drop entirely",
                            },
                        },
                    },
                },
                "required": ["dataset_id", "columns"],
            },
        ),
        Tool(
            name="run_clustering",
            description=(
                "对选定列执行聚类。使用HDBSCAN（默认）或KMeans。特征在聚类前进行"
                "标准化缩放。返回聚类标签和质量指标。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Dataset identifier",
                    },
                    "cluster_columns": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Columns to use as clustering features",
                    },
                    "algorithm": {
                        "type": "string",
                        "enum": ["hdbscan", "kmeans"],
                        "description": "Clustering algorithm",
                        "default": "hdbscan",
                    },
                    "params": {
                        "type": "object",
                        "description": (
                            "Algorithm parameters. HDBSCAN: min_cluster_size (5), "
                            "min_samples (None), metric ('euclidean'). "
                            "KMeans: n_clusters (3), random_state (42)."
                        ),
                    },
                    "random_state": {
                        "type": "integer",
                        "description": "Random seed for reproducibility",
                        "default": 42,
                    },
                },
                "required": ["dataset_id", "cluster_columns"],
            },
        ),
        Tool(
            name="evaluate_clustering",
            description=(
                "评估数据集的聚类结果。当前支持轮廓系数（silhouette）、Davies-Bouldin指数、"
                "Calinski-Harabasz指数、噪声比（noise_ratio）和簇大小（cluster_sizes）指标。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "cluster_result_id": {
                        "type": "string",
                        "description": "Result ID from run_clustering",
                    },
                    "dataset_id": {
                        "type": "string",
                        "description": "Original dataset for metric computation",
                    },
                    "metrics": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": (
                            "Metrics to compute. Options: silhouette, "
                            "davies_bouldin, calinski_harabasz, noise_ratio, "
                            "cluster_sizes (default: all)"
                        ),
                    },
                },
                "required": ["cluster_result_id", "dataset_id"],
            },
        ),
        Tool(
            name="extract_features",
            description=(
                "提取每个聚类的区分特征并自动保存到Django ORM（ClusterFeature模型）。"
                "使用z-score或ANOVA对每个聚类的特征重要性进行排序。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Dataset containing the feature columns",
                    },
                    "cluster_result_id": {
                        "type": "string",
                        "description": "Result ID from run_clustering",
                    },
                    "top_k": {
                        "type": "integer",
                        "description": "Top distinguishing features per cluster",
                        "default": 10,
                    },
                    "method": {
                        "type": "string",
                        "enum": ["zscore", "anova"],
                        "description": "Method to rank feature importance",
                        "default": "zscore",
                    },
                    "save_to_db": {
                        "type": "boolean",
                        "description": "Save features to Django ORM if True",
                        "default": True,
                    },
                },
                "required": ["dataset_id", "cluster_result_id"],
            },
        ),
        Tool(
            name="export_results",
            description=(
                "将数据集或聚类结果以CSV或Parquet格式导出到磁盘。数据集物化后写入；"
                "聚类结果序列化为JSON。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "result_id": {
                        "type": "string",
                        "description": "Dataset ID or cluster result ID to export",
                    },
                    "output_path": {
                        "type": "string",
                        "description": "Output file path or directory",
                    },
                    "format": {
                        "type": "string",
                        "enum": ["csv", "parquet", "json"],
                        "description": "Output format",
                        "default": "csv",
                    },
                    "overwrite": {
                        "type": "boolean",
                        "description": "Overwrite existing file",
                        "default": False,
                    },
                },
                "required": ["result_id", "output_path"],
            },
        ),
        Tool(
            name="list_datasets",
            description="列出会话存储中的所有活跃数据集和结果。",
            inputSchema={
                "type": "object",
                "properties": {},
            },
        ),
        Tool(
            name="drop_dataset",
            description=(
                "从会话存储中移除数据集（或结果）并通过垃圾回收释放内存。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Dataset or result identifier to remove",
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="clone_dataset",
            description=(
                "浅拷贝数据集。由于LazyFrame是查询计划，不会复制实际数据。适用于分支分析。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Source dataset identifier",
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="filter_and_cluster",
            description=(
                "一步完成：先应用列级过滤器，再对过滤结果执行聚类。支持eq/neq/gt/gte/lt/lte/"
                "contains/not_contains/in/not_in/is_null/is_not_null运算符和HDBSCAN/KMeans聚类。"
                "自动选择数值列作为聚类特征。将过滤后的数据集和聚类结果都存储在会话中。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Source dataset identifier",
                    },
                    "filters": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "column": {"type": "string"},
                                "op": {
                                    "type": "string",
                                    "enum": [
                                        "eq", "neq", "gt", "gte", "lt", "lte",
                                        "contains", "not_contains",
                                        "in", "not_in",
                                        "is_null", "is_not_null",
                                    ],
                                },
                                "value": {"description": "Operand value"},
                            },
                            "required": ["column", "op"],
                        },
                        "description": "List of filter expressions",
                    },
                    "logic": {
                        "type": "string",
                        "enum": ["and", "or"],
                        "description": "Combine filters with AND or OR logic",
                        "default": "and",
                    },
                    "algorithm": {
                        "type": "string",
                        "enum": ["hdbscan", "kmeans"],
                        "description": "Clustering algorithm",
                        "default": "hdbscan",
                    },
                    "params": {
                        "type": "object",
                        "description": (
                            "Algorithm parameters. HDBSCAN: min_cluster_size (5), "
                            "min_samples (None), metric ('euclidean'). "
                            "KMeans: n_clusters (3), random_state (42)."
                        ),
                    },
                },
                "required": ["dataset_id", "filters"],
            },
        ),
        Tool(
            name="build_entity_profiles",
            description=(
                "按实体列（源IP、SNI等）分组原始流数据。自动检测实体列（如 src_ip, "
                "dst_ip, sni, host, mac, domain），按列分组并计算每个实体的聚合特征："
                "流数量、唯一端口数、TLS版本分布、字节总量、唯一目标数。"
                "生成的数据集每个实体一行。结果存储在会话中，可供后续 compute_scores 使用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Dataset ID from load_data or profile_data",
                    },
                    "entity_column": {
                        "type": "string",
                        "description": "Specific entity column to group by (e.g. 'src_ip')",
                    },
                    "entity_columns": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Multiple entity columns to group by",
                    },
                    "auto_detect": {
                        "type": "boolean",
                        "description": "Auto-detect entity columns (IP, SNI, host, MAC, domain)",
                        "default": True,
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="compute_scores",
            description=(
                "在实体数据上计算自适应代理/异常/威胁评分。从数据自动学习"
                "特征权重，计算加权代理评分（proxy_score）、基于IQR的风险等级"
                "（normal/watch/suspicious/critical）和威胁评分（threat_score）。"
                "返回评分摘要与阈值。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Entity dataset ID from filter_and_cluster or a scored dataset",
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="detect_anomalies",
            description=(
                "对已评分的实体数据运行孤立森林（Isolation Forest）异常检测。识别偏离"
                "正常模式的实体。分块处理以应对大数据集。返回异常实体ID列表。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Scored entity dataset ID (after compute_scores)",
                    },
                    "contamination": {
                        "type": "number",
                        "description": "Expected proportion of anomalies (0.0-0.5)",
                        "default": 0.05,
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="visualize_anomalies",
            description=(
                "为高风险实体生成UMAP-2D嵌入和聚类可视化数据。按风险等级和聚类着色。"
                "返回用于前端渲染的散点图JSON数据。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Scored entity dataset ID (after compute_scores)",
                    },
                },
                "required": ["dataset_id"],
            },
        ),
        # ── Edge-case / Diagnostic tools (unexpected situations) ────────
        Tool(
            name="validate_data",
            description=(
                "诊断工具：检查数据质量问题——文件间schema一致性、缺失值比例、列类型冲突、"
                "编码问题和空列。当流水线步骤因schema错误失败，或怀疑CSV数据损坏/格式异常时使用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
                    "fix_issues": {"type": "boolean", "description": "Auto-fix schema inconsistencies", "default": False},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="explore_distributions",
            description=(
                "诊断工具：计算每列分布统计（唯一值计数、空值比例、最小/最大/均值/标准差、"
                "前10高频值、直方图分箱）。当聚类产生意外结果、实体检测失败，或需要在选择"
                "分析策略前了解数据时使用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data or entity data"},
                    "columns": {"type": "array", "items": {"type": "string"},
                                "description": "Specific columns to analyze (default: auto-detect numeric)"},
                    "max_sample": {"type": "integer", "description": "Max rows to sample", "default": 10000},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="find_outliers",
            description=(
                "诊断工具：使用IQR方法检测数值列中的统计异常值。返回任何数值列超过中位数3倍IQR的行。"
                "当怀疑数据质量问题或极端值偏离分析时使用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
                    "columns": {"type": "array", "items": {"type": "string"},
                                "description": "Columns to check (default: all numeric)"},
                    "threshold": {"type": "number", "description": "IQR multiplier (default 3.0)", "default": 3.0},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="diagnose_clustering",
            description=(
                "诊断工具：当聚类产生太少/太多聚类、全部为噪声或轮廓系数低时——诊断原因。"
                "检查：特征方差、相关矩阵、特征降维分析、数据稀疏性，并建议参数调整"
                "（min_cluster_size、算法选择、特征选择）。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Entity dataset ID"},
                    "cluster_result_id": {"type": "string", "description": "Cluster result ID from previous run_clustering"},
                },
                "required": ["dataset_id", "cluster_result_id"],
            },
        ),
        Tool(
            name="compare_datasets",
            description=(
                "诊断工具：比较两个数据集（如处理前/后，或两次不同运行）。报告schema差异、"
                "行数变化、列值分布偏移。用于验证处理步骤未损坏数据或比较不同运行的结果。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id_a": {"type": "string", "description": "First dataset ID"},
                    "dataset_id_b": {"type": "string", "description": "Second dataset ID"},
                },
                "required": ["dataset_id_a", "dataset_id_b"],
            },
        ),
        Tool(
            name="export_debug_sample",
            description=(
                "诊断工具：导出原始数据样本为JSON用于外部调试。当流水线深处发生错误且需要"
                "检查实际数据值时使用。返回前N行所有列。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID"},
                    "max_rows": {"type": "integer", "description": "Rows to export", "default": 50},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="repair_schema",
            description=(
                "诊断工具：通过对齐列名（不区分大小写合并、下划线/连字符规范化）并用空值填充"
                "缺失列，尝试修复文件间的schema不匹配。当load_data因CSV文件具有不同列集而"
                "失败时使用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Failed dataset ID from load_data"},
                },
                "required": ["dataset_id"],
            },
        ),
        # ── Read-only analysis tools (LLM deep-dive, no side effects) ────
        Tool(
            name="analyze_patterns",
            description=(
                "分析工具：检测流量模式——top源/目标IP、端口分布、TLS版本分布、协议混合。"
                "返回汇总统计。只读，不修改数据。用于了解高层流量结构。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
                    "top_n": {"type": "integer", "description": "Number of top items per category", "default": 10},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="analyze_temporal",
            description=(
                "分析工具：分析随时间变化的流量——小时/日流量计数、繁忙时段、空闲时段。"
                "需要时间戳列。只读。用于了解时间相关模式和周期性。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
                    "timestamp_column": {"type": "string", "description": "Column name for timestamps (auto-detected if omitted)"},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="analyze_fft",
            description=(
                "FFT频谱分析——从时间戳列提取周期模式和频率特征。适用于检测时间周期性"
                "（如每日/每周流量模式）。返回频率、幅度、top周期、频谱质心、周期性评分。"
                "只读，无副作用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Dataset ID from load_data",
                    },
                    "timestamp_column": {
                        "type": "string",
                        "description": "Column name containing timestamps or numeric time values",
                    },
                    "top_n": {
                        "type": "integer",
                        "description": "Number of top frequency peaks to return",
                        "default": 5,
                    },
                },
                "required": ["dataset_id", "timestamp_column"],
            },
        ),
        Tool(
            name="analyze_tls_health",
            description=(
                "分析工具：TLS安全态势评估——检查过时的TLS版本（1.0/1.1）、弱加密套件、"
                "证书问题、SNI异常。只读。用于评估流量的TLS安全性。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="analyze_geo_distribution",
            description=(
                "分析工具：流量地理分布——top源/目标国家、异常地点对（如不可能的行程）、"
                "ISP多样性。只读。需要国家或经纬度列。用于检测地理异常。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="analyze_entity_detail",
            description=(
                "分析工具：深入分析单个实体——其所有流、使用的TLS版本、目标多样性、时间模式、"
                "异常评分。只读。在compute_scores后调查特定可疑实体时使用。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "Entity dataset ID (after filter_and_cluster or compute_scores)"},
                    "entity_value": {"type": "string", "description": "Entity value to investigate (e.g. an IP address or SNI)"},
                },
                "required": ["dataset_id", "entity_value"],
            },
        ),
        Tool(
            name="compute_distance_matrix",
            description=(
                "LLM驱动工具：对数据集的每一行执行用户编写的Python函数，返回距离/相似度评分摘要。"
                "LLM提供函数体，签名格式为`def distance_fn(row: dict) -> float:`。执行器将其包装，"
                "在受限命名空间（仅math、numpy）中逐行运行，返回得分的最小/最大/均值/标准差以及"
                "评分行样本。用于现有工具无法表达的一次性度量、过滤或评分任务。"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {
                        "type": "string",
                        "description": "Source dataset identifier",
                    },
                    "python_function": {
                        "type": "string",
                        "description": (
                            "Body of a Python function with signature `def distance_fn(row: dict) -> float:`. "
                            "Example: \"return abs(row.get('col1', 0)) + row.get('col2', 0)\". "
                            "Available namespace: math, numpy (as np). No dangerous modules."
                        ),
                    },
                    "columns": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Columns to pass to the function (default: auto-detect numeric columns)",
                    },
                    "max_sample": {
                        "type": "integer",
                        "description": "Maximum rows to evaluate (default 10000 for performance)",
                        "default": 10000,
                    },
                },
                "required": ["dataset_id", "python_function"],
            },
        ),
    ]


# ═══════════════════════════════════════════════════════════════════════
#  Handler implementations
# ═══════════════════════════════════════════════════════════════════════

async def handle_call(name: str, arguments: dict) -> dict:
    """Dispatch a tool call to the appropriate handler.

    Args:
        name: Tool name (must match one of the 11 tools).
        arguments: Tool-specific parameters.

    Returns:
        A JSON-serialisable dict.

    Raises:
        ValueError: If *name* is not a recognised tool.
    """
    _handlers = {
        'load_data': _handle_load_data,
        'profile_data': _handle_profile_data,
        'filter_data': _handle_filter_data,
        'preprocess_data': _handle_preprocess_data,
        'run_clustering': _handle_run_clustering,
        'evaluate_clustering': _handle_evaluate_clustering,
        'extract_features': _handle_extract_features,
        'export_results': _handle_export_results,
        'list_datasets': _handle_list_datasets,
        'drop_dataset': _handle_drop_dataset,
        'clone_dataset': _handle_clone_dataset,
        'build_entity_profiles': _handle_build_entity_profiles,
        'compute_scores': _handle_compute_scores,
        'filter_and_cluster': _handle_filter_and_cluster,
        'detect_anomalies': _handle_detect_anomalies,
        'visualize_anomalies': _handle_visualize_anomalies,
        'validate_data': _handle_validate_data,
        'explore_distributions': _handle_explore_distributions,
        'find_outliers': _handle_find_outliers,
        'diagnose_clustering': _handle_diagnose_clustering,
        'compare_datasets': _handle_compare_datasets,
        'export_debug_sample': _handle_export_debug_sample,
        'repair_schema': _handle_repair_schema,
    'analyze_patterns': _handle_analyze_patterns,
                'analyze_temporal': _handle_analyze_temporal,
                'analyze_fft': _handle_analyze_fft,
                'analyze_tls_health': _handle_analyze_tls_health,
        'analyze_geo_distribution': _handle_analyze_geo_distribution,
        'analyze_entity_detail': _handle_analyze_entity_detail,
        'compute_distance_matrix': _handle_compute_distance_matrix,
    }
    handler = _handlers.get(name)
    if handler is None:
        raise ValueError(f"Unknown tool: '{name}'")
    result = await handler(**arguments)
    # Always ensure truncated key exists
    if 'truncated' not in result:
        result['truncated'] = False
    return _truncate_response(result)


# ── 1. load_data ────────────────────────────────────────────────────────

async def _handle_load_data(
    csv_glob: str,
    config_path: Optional[str] = None,
    encoding: str = 'utf-8',
    delimiter: str = ',',
    schema_strict: bool = False,
    recursive: bool = False,
) -> dict:
    """Load, validate, merge CSVs, and store in session."""
    try:
        lf, schema, row_count, file_count, memory_mb = load_csv_directory(
            glob_pattern=csv_glob,
            encoding=encoding,
            delimiter=delimiter,
            config_path=config_path,
            sample_rows=10000,
            schema_strict=schema_strict,
            recursive=recursive,
        )
    except FileNotFoundError as e:
        return {'error': f'File not found: {e}', 'truncated': False}
    except ValueError as e:
        return {'error': f'Schema validation failed: {e}', 'truncated': False}

    store = SessionStore()
    dataset_id = _generate_id('ds')
    store.store_dataset(
        dataset_id=dataset_id,
        lazyframe=lf,
        schema=schema,
        metadata={
            'row_count': row_count,
            'file_count': file_count,
            'memory_mb': round(memory_mb, 2),
            'csv_glob': csv_glob,
            'encoding': encoding,
        },
    )

    return {
        'dataset_id': dataset_id,
        'schema': schema,
        'row_count': row_count,
        'file_count': file_count,
        'memory_mb': round(memory_mb, 2),
    }


# ── 2. profile_data ─────────────────────────────────────────────────────

async def _handle_profile_data(
    dataset_id: str,
    sample_size: int = 1000,
    columns: Optional[list[str]] = None,
) -> dict:
    """Profile dataset columns and correlations."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf: pl.LazyFrame = entry['lazyframe']
    profile = _profile_dataset(lf, sample_size=sample_size, columns=columns)
    return profile


# ── 3. filter_data ──────────────────────────────────────────────────────

def _build_filter_expr(column: str, op: str, value: Any) -> pl.Expr:
    """Build a single Polars expression from a filter descriptor."""
    col = pl.col(column)
    op_map = {
        'eq': col == value,
        'neq': col != value,
        'gt': col > value,
        'gte': col >= value,
        'lt': col < value,
        'lte': col <= value,
        'contains': col.str.contains(value) if value is not None else col.cast(str).str.contains(''),
        'not_contains': ~(col.str.contains(value) if value is not None else col.cast(str).str.contains('')),
        'in': col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value]),
        'not_in': ~(col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value])),
        'is_null': col.is_null(),
        'is_not_null': col.is_not_null(),
    }
    expr = op_map.get(op)
    if expr is None:
        raise ValueError(f"Unsupported filter operator: {op}")
    return expr


async def _handle_filter_data(
    dataset_id: str,
    filters: list[dict],
    logic: str = 'and',
) -> dict:
    """Apply filters and store filtered dataset."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

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

    exprs = [_build_filter_expr(
        f['column'], f.get('op', 'eq'), f.get('value'),
    ) for f in filters]

    if not exprs:
        return {'error': 'No filter expressions could be built', 'truncated': False}

    if logic == 'and':
        combined = exprs[0]
        for e in exprs[1:]:
            combined = combined & e
    else:
        combined = exprs[0]
        for e in exprs[1:]:
            combined = combined | e

    filtered_lf = lf.filter(combined)

    # Estimate row count (fast streaming count)
    try:
        row_count = filtered_lf.select(pl.lit(1).count()).collect(streaming=True).item()
    except Exception:
        row_count = None

    new_id = _generate_id('flt')
    store.store_dataset(
        dataset_id=new_id,
        lazyframe=filtered_lf,
        schema=full_schema,
        metadata={
            'row_count': row_count,
            'parent_dataset_id': dataset_id,
            'filter_count': len(filters),
            'logic': logic,
        },
        parent_id=dataset_id,
    )

    return {
        'filtered_dataset_id': new_id,
        'row_count': row_count,
        'column_deltas': {'kept': list(full_schema.keys()), 'dropped': []},
    }


# ── 4. preprocess_data ──────────────────────────────────────────────────

async def _handle_preprocess_data(
    dataset_id: str,
    columns: list[str],
    config: Optional[dict] = None,
) -> dict:
    """Preprocess selected columns (scale, encode, impute, drop)."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf: pl.LazyFrame = entry['lazyframe']
    cfg = config or {}
    column_mapping: dict[str, str] = {}

    df = lf.collect(streaming=True)

    # Resolve columns that actually exist
    available = set(df.columns)
    target_cols = [c for c in columns if c in available]
    if not target_cols:
        return {'error': 'None of the specified columns exist in the dataset',
                'truncated': False}

    # Drop columns
    drop_cols = cfg.get('drop', [])
    if drop_cols:
        df = df.drop([c for c in drop_cols if c in df.columns])
        for c in drop_cols:
            if c in available:
                column_mapping[c] = f'<dropped:{c}>'

    # Fill missing values
    fill_strategy = cfg.get('fillna')
    if fill_strategy and fill_strategy != 'drop':
        for col in target_cols:
            if col not in df.columns:
                continue
            if df[col].null_count() == 0:
                continue
            if fill_strategy == 'mean':
                mean_val = df[col].drop_nulls().mean()
                if mean_val is not None:
                    df = df.with_columns(pl.col(col).fill_null(mean_val))
            elif fill_strategy == 'median':
                med_val = df[col].drop_nulls().median()
                if med_val is not None:
                    df = df.with_columns(pl.col(col).fill_null(med_val))
            elif fill_strategy == 'zero':
                df = df.with_columns(pl.col(col).fill_null(0))
    elif fill_strategy == 'drop':
        df = df.drop_nulls(subset=target_cols)

    # Standard scaling (z-score)
    if cfg.get('standardize'):
        from sklearn.preprocessing import StandardScaler
        num_cols = [c for c in target_cols if c in df.columns
                    and df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64)]
        if num_cols:
            scaler = StandardScaler()
            scaled = scaler.fit_transform(df.select(num_cols).to_numpy())
            for i, col in enumerate(num_cols):
                col_scaled = f'{col}_scaled'
                df = df.with_columns(pl.Series(col_scaled, scaled[:, i]))
                column_mapping[col] = col_scaled

    # One-hot encoding
    if cfg.get('onehot'):
        str_cols = [c for c in target_cols if c in df.columns
                    and df[c].dtype in (pl.Utf8, pl.String, pl.Categorical)]
        if str_cols:
            df = df.to_dummies(columns=str_cols, drop_first=False)
            for c in str_cols:
                column_mapping[c] = f'<onehot:{c}>'

    new_id = _generate_id('pp')
    new_lf = df.lazy()
    new_schema = {name: str(dtype) for name, dtype in
                  zip(new_lf.collect_schema().names(), new_lf.collect_schema().dtypes())}

    store.store_dataset(
        dataset_id=new_id,
        lazyframe=new_lf,
        schema=new_schema,
        metadata={
            'row_count': len(df),
            'parent_dataset_id': dataset_id,
            'preprocessing_config': cfg,
        },
        parent_id=dataset_id,
    )

    return {
        'preprocessed_dataset_id': new_id,
        'column_mapping': column_mapping,
        'row_count': len(df),
        'column_count': len(new_schema),
    }


# ── 5. run_clustering ───────────────────────────────────────────────────

async def _handle_run_clustering(
    dataset_id: str,
    cluster_columns: list[str],
    algorithm: str = 'hdbscan',
    params: Optional[dict] = None,
    random_state: int = 42,
) -> dict:
    """Run HDBSCAN or KMeans clustering on selected columns (async)."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf: pl.LazyFrame = entry['lazyframe']
    p = params or {}

    try:
        # Collect cluster features - numeric only
        available = lf.collect_schema()
        available_names = available.names()
        available_dtypes = available.dtypes()
        numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                          pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                          pl.Float32, pl.Float64)
        feature_cols = [
            c for c in cluster_columns
            if c in available_names
            and available_dtypes[available_names.index(c)] in numeric_dtypes
        ]
        if not feature_cols:
            # Fall back: auto-select all numeric columns (exclude internal ones)
            feature_cols = [
                available_names[i] for i, dt in enumerate(available_dtypes)
                if dt in numeric_dtypes and not available_names[i].startswith('_')
            ][:10]
        if not feature_cols:
            # Fallback: try coercing Utf8 columns to Float64 (common when data loaded from
            # SQLite TEXT storage without proper dtype restoration).
            utf8_cols = [
                available_names[i] for i, dt in enumerate(available_dtypes)
                if dt == pl.Utf8 and not available_names[i].startswith('_')
            ][:10]
            if utf8_cols:
                # Attempt per-column cast; filter to columns where coercion succeeded
                for col in utf8_cols:
                    try:
                        casted = lf.select(
                            pl.col(col).cast(pl.Float64, strict=False).alias(col)
                        ).collect(streaming=True)
                        null_ratio = casted[col].is_null().sum() / max(len(casted), 1)
                        if null_ratio < 0.9:  # at least 10% successfully parsed as floats
                            lf = lf.with_columns(
                                pl.col(col).cast(pl.Float64, strict=False)
                            )
                            feature_cols.append(col)
                    except Exception:
                        pass
            if not feature_cols:
                all_numeric = [f"{available_names[i]}" for i, dt in enumerate(available_dtypes)
                               if dt in numeric_dtypes]
                return {
                    'error': f'No numeric cluster columns found in dataset. Available: {all_numeric}',
                    'truncated': False
                }

        # Downsample before collect: compute row count via fast lazy count
        try:
            row_count = lf.select(pl.len()).collect(streaming=True).item()
        except Exception:
            row_count = 0
        if row_count > 100000:
            # Sample to 100K before collect to avoid OOM on 10M+ rows
            lf = lf.collect(streaming=True).sample(n=100000, seed=42).lazy()

        df_features = lf.select(feature_cols).collect(streaming=True)

        # Drop all-NaN columns (keep rows, drop useless columns)
        for col in df_features.columns:
            if df_features[col].is_null().all():
                df_features = df_features.drop(col)
                feature_cols = [c for c in feature_cols if c != col]

        if df_features.width == 0:
            return {'error': 'No valid feature columns after dropping all-NaN columns',
                    'truncated': False}

        data = df_features.to_numpy()

        # Fill remaining NaN with column mean
        if data.dtype.kind == 'f':
            col_mean = np.nanmean(data, axis=0)
            col_mean = np.nan_to_num(col_mean, nan=0.0)
            data = np.where(np.isnan(data), col_mean, data)

        # Zero-sample / too-few-samples guard
        if len(data) < 3:
            return {
                'cluster_result_id': _generate_id('clust'),
                'n_clusters': 0,
                'n_noise': 0,
                'quality_metrics': {'note': 'too few samples'},
                'truncated': False,
            }

        # ── D14: Low memory check (safe fallback if psutil not installed) ──
        low_memory = False
        try:
            import psutil
            mem = psutil.virtual_memory()
            low_memory = mem.available < 2 * 1024**3  # 2 GB
        except ImportError:
            pass

        # ── D14: Auto-downsample for low memory (>100K rows → 50K) ────────
        if low_memory and len(data) > 100000:
            rng = np.random.default_rng(42)
            idx = rng.choice(len(data), 50000, replace=False)
            data = data[idx]

        # ── D5: Downsampling → extract matching indices ──────────────
        # Downsample to 50K max, store indices so _handle_extract_features
        # can filter the same rows.
        _ds_idx = None
        if len(data) > 50000:
            rng = np.random.default_rng(42)
            _ds_idx = rng.choice(len(data), 50000, replace=False)
            data = data[_ds_idx]

        # ── H17: Variance filtering ───────────────────────────────────────
        # Remove features with zero / near-zero variance (< 1e-10)
        if data.shape[1] > 1:
            variances = np.var(data, axis=0)
            keep_var = np.where(variances >= 1e-10)[0]
            if len(keep_var) == 0:
                # All features have zero variance — fall back to unfiltered data
                keep_var = np.arange(data.shape[1])
            if len(keep_var) < data.shape[1]:
                data = data[:, keep_var]
                feature_cols = [feature_cols[i] for i in keep_var]

        # Log clustering input
        import logging
        _clust_logger = logging.getLogger(__name__)
        _clust_logger.info(f'[CLUSTER_INPUT] rows={len(data)} cols={data.shape[1] if len(data.shape) > 1 else 1} feature_cols={feature_cols}')

        # Standard scale
        from sklearn.preprocessing import StandardScaler
        _clust_logger.info('[CLUSTER_SCALE] Starting StandardScaler...')
        data_scaled = StandardScaler().fit_transform(data)
        _clust_logger.info(f'[CLUSTER_SCALE] Done. shape={data_scaled.shape}')

        # ── H17: Correlation filtering ────────────────────────────────────
        # Remove features with Pearson correlation > 0.95 (keep the first one)
        if data_scaled.shape[1] > 1:
            corr = np.corrcoef(data_scaled.T)
            upper = np.triu(np.abs(corr), 1)
            to_drop = [j for j in range(upper.shape[1]) if any(upper[:, j] > 0.95)]
            if to_drop:
                keep = [j for j in range(data_scaled.shape[1]) if j not in to_drop]
                data_scaled = data_scaled[:, keep]
                feature_cols = [feature_cols[j] for j in keep]

        # ── D14: SVD dimensionality reduction (if >50 features) ───────────
        n_features = data_scaled.shape[1]
        if n_features > 50:
            from sklearn.decomposition import TruncatedSVD
            svd = TruncatedSVD(n_components=50, random_state=random_state)
            data_scaled = svd.fit_transform(data_scaled)

        # ── H17: HDBSCAN hyperparameter auto-tuning based on data size ────
        n = len(data_scaled)
        if algorithm == 'hdbscan':
            min_cluster = p.get('min_cluster_size', max(3, min(50, n // 20)))
            p['min_cluster_size'] = min_cluster
            p['min_samples'] = p.get('min_samples', min_cluster)
            p['cluster_selection_epsilon'] = p.get('cluster_selection_epsilon', 0.0)
        else:
            # KMeans safety cap
            p['min_cluster_size'] = min(
                p.get('min_cluster_size', 5),
                max(n // 2, 2),
            )

        # Clustering
        from sklearn.cluster import HDBSCAN, MiniBatchKMeans
        if algorithm == 'kmeans':
            n_clusters = p.get('n_clusters', 3)
            model = MiniBatchKMeans(
                n_clusters=n_clusters, random_state=random_state,
                batch_size=1024, n_init='auto',
            )
        else:
            model = HDBSCAN(
                min_cluster_size=p['min_cluster_size'],
                min_samples=p.get('min_samples', None),
                metric=p.get('metric', 'euclidean'),
                cluster_selection_epsilon=p.get('cluster_selection_epsilon', 0.0),
            )

        labels_array = model.fit_predict(data_scaled)
        n_clusters_found = int(len(set(labels_array)) - (1 if -1 in labels_array else 0))
        noise_count = int((labels_array == -1).sum()) if algorithm == 'hdbscan' else 0
        noise_ratio = round(noise_count / len(labels_array), 4) if algorithm == 'hdbscan' else 0.0

        # Map back to original rows
        full_labels = labels_array

        # Quality metrics
        quality = {
            'n_clusters': n_clusters_found,
            'n_noise': noise_count,
            'noise_ratio': noise_ratio,
            'algorithm': algorithm,
        }

        if n_clusters_found > 1:
            try:
                from sklearn.metrics import silhouette_score
                score = silhouette_score(data_scaled, labels_array, random_state=random_state)
                quality['silhouette_score'] = round(float(score), 4)
            except Exception:
                quality['silhouette_score'] = None
            try:
                from sklearn.metrics import davies_bouldin_score
                dbs = davies_bouldin_score(data_scaled, labels_array)
                quality['davies_bouldin_score'] = round(float(dbs), 4)
            except Exception:
                quality['davies_bouldin_score'] = None
            try:
                from sklearn.metrics import calinski_harabasz_score
                chs = calinski_harabasz_score(data_scaled, labels_array)
                quality['calinski_harabasz_score'] = round(float(chs), 4)
            except Exception:
                quality['calinski_harabasz_score'] = None

        # Cluster sizes
        cluster_sizes = {}
        for label in np.unique(labels_array):
            cluster_sizes[int(label)] = int((labels_array == label).sum())
        quality['cluster_sizes'] = cluster_sizes

    except Exception:
        raise

    result_id = _generate_id('clust')
    store.store_cluster_result(
        result_id=result_id,
        labels=full_labels.tolist(),
        n_clusters=n_clusters_found,
        params={
            'algorithm': algorithm,
            'params': p,
            'random_state': random_state,
            'feature_columns': feature_cols,
            '_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None,
        },
        parent_dataset_id=dataset_id,
    )

    return {
        'cluster_result_id': result_id,
        'n_clusters': n_clusters_found,
        'n_noise': noise_count,
        'labels_sample': full_labels[:50].tolist(),
        'quality_metrics': quality,
    }


# ── 6. evaluate_clustering ──────────────────────────────────────────────

async def _handle_evaluate_clustering(
    cluster_result_id: str,
    dataset_id: str,
    metrics: Optional[list[str]] = None,
) -> dict:
    """Evaluate clustering result against the dataset."""
    store = SessionStore()
    cluster_entry = store.get_cluster_result(cluster_result_id)
    if cluster_entry is None:
        return {'error': f'Cluster result not found: {cluster_result_id}',
                'truncated': False}

    dataset_entry = store.get_dataset(dataset_id)
    if dataset_entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    labels = cluster_entry['labels']
    default_metrics = ['silhouette', 'davies_bouldin', 'calinski_harabasz',
                       'noise_ratio', 'cluster_sizes']
    requested = metrics or default_metrics

    scores: dict = {'n_clusters': cluster_entry['n_clusters']}

    try:
        lf: pl.LazyFrame = dataset_entry['lazyframe']
        params = cluster_entry.get('params', {})
        feature_cols = params.get('feature_columns', [])
        if not feature_cols and 'feature_columns' in params:
            pass
        elif feature_cols:
            df_features = lf.select(feature_cols).collect(streaming=True)
            data = df_features.to_numpy()

            # Align labels with data (handle missing rows)
            valid_mask = ~np.isnan(data).any(axis=1) if data.dtype.kind == 'f' else np.ones(len(data), dtype=bool)
            data_clean = data[valid_mask]
            labels_clean = np.array(labels)[:len(data_clean)]

            unique_labels = set(labels_clean)
            n_clusters_actual = len(unique_labels - {-1})

            for m in requested:
                if m == 'silhouette' and n_clusters_actual >= 2:
                    from sklearn.metrics import silhouette_score
                    scores['silhouette'] = round(float(
                        silhouette_score(data_clean, labels_clean, random_state=42)), 4)
                elif m == 'davies_bouldin' and n_clusters_actual >= 2:
                    from sklearn.metrics import davies_bouldin_score
                    scores['davies_bouldin'] = round(float(
                        davies_bouldin_score(data_clean, labels_clean)), 4)
                elif m == 'calinski_harabasz' and n_clusters_actual >= 2:
                    from sklearn.metrics import calinski_harabasz_score
                    scores['calinski_harabasz'] = round(float(
                        calinski_harabasz_score(data_clean, labels_clean)), 4)
                elif m == 'noise_ratio':
                    noise = int((labels_clean == -1).sum())
                    scores['noise_ratio'] = round(noise / len(labels_clean), 4)
                elif m == 'cluster_sizes':
                    sizes = {}
                    for label in unique_labels:
                        sizes[int(label)] = int((labels_clean == label).sum())
                    scores['cluster_sizes'] = sizes

        scores['n_clusters'] = cluster_entry['n_clusters']

    except Exception:
        raise

    return {'scores': scores}


# ── 7. extract_features ─────────────────────────────────────────────────

async def _handle_extract_features(
    dataset_id: str,
    cluster_result_id: str,
    top_k: int = 10,
    method: str = 'zscore',
    save_to_db: bool = True,
    run_id: int | None = None,
) -> dict:
    """Compute per-cluster distinguishing features and optionally persist."""
    store = SessionStore()
    dataset_entry = store.get_dataset(dataset_id)
    if dataset_entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    cluster_entry = store.get_cluster_result(cluster_result_id)
    if cluster_entry is None:
        return {'error': f'Cluster result not found: {cluster_result_id}',
                'truncated': False}

    try:
        lf: pl.LazyFrame = dataset_entry['lazyframe']
        labels = cluster_entry['labels']

        # D5: If clustering downsampled the data, filter matching rows
        _ds_idx = cluster_entry.get('params', {}).get('_ds_idx')
        if _ds_idx is not None:
            # Filter in lazy plan before collect — avoids loading full 10M rows
            df = (lf.with_row_index('__ds_idx')
                  .filter(pl.col('__ds_idx').is_in(_ds_idx))
                  .drop('__ds_idx')
                  .collect(streaming=True))
        else:
            df = lf.collect(streaming=True)

        # Truncate labels if df was filtered but labels weren't
        if len(labels) != len(df):
            labels = labels[:len(df)]

        # Only use numeric columns for feature extraction
        numeric_cols = [c for c in df.columns
                        if df[c].dtype in (pl.Float32, pl.Float64,
                                            pl.Int32, pl.Int64,
                                            pl.UInt32, pl.UInt64)]
        if len(labels) != len(df):
            labels = labels[:len(df)]

        # Global statistics per column
        global_stats = {}
        for col in numeric_cols:
            s = df[col].drop_nulls()
            if len(s) > 0:
                global_stats[col] = {
                    'mean': float(s.mean()),
                    'std': float(s.std()) if s.std() is not None else 1.0,
                }

        unique_labels = sorted(set(labels))
        all_features = []

        for label in unique_labels:
            mask = np.array(labels) == label
            if mask.sum() < 2:
                continue

            cluster_df = df.filter(pl.Series('__mask__', mask))
            cluster_features = []

            for col in numeric_cols:
                c_series = cluster_df[col].drop_nulls()
                gs = global_stats.get(col)
                if gs is None or gs['std'] == 0 or len(c_series) == 0:
                    continue

                c_mean = float(c_series.mean())
                c_std = float(c_series.std()) if c_series.std() is not None else 0.0

                if method == 'zscore':
                    score = (c_mean - gs['mean']) / max(gs['std'], 1e-10)
                else:
                    # ANOVA-like: F = (between-group variance) / (within-group variance)
                    score = abs(c_mean - gs['mean']) / max(c_std + gs['std'], 1e-10)

                cluster_features.append({
                    'feature_name': col,
                    'cluster_label': int(label),
                    'mean': round(c_mean, 4),
                    'std': round(c_std, 4),
                    'global_mean': round(gs['mean'], 4),
                    'global_std': round(gs['std'], 4),
                    'distinguishing_score': round(float(score), 4),
                    'distinguishing_method': method,
                })

            # Sort by distinguishing score descending, take top_k
            cluster_features.sort(key=lambda x: abs(x['distinguishing_score']), reverse=True)
            all_features.extend(cluster_features[:top_k])

        # ── Compute UMAP-2D embedding ──────────────────────────────────────
        if len(numeric_cols) >= 2 and len(df) > 2:
            try:
                from sklearn.preprocessing import StandardScaler
                import umap
                MAX_UMAP_TRAIN = 10000
                BATCH_SIZE = 1000
                data_matrix = df.select(numeric_cols).to_numpy()
                if len(df) > MAX_UMAP_TRAIN:
                    # Train UMAP on 10K sample, batch-transform full dataset
                    idx_sample = np.random.RandomState(42).choice(
                        len(df), size=MAX_UMAP_TRAIN, replace=False)
                    mat_sample = data_matrix[idx_sample]
                    scaler = StandardScaler().fit(mat_sample)
                    mat_sample_scaled = np.nan_to_num(scaler.transform(mat_sample),
                                                      nan=0.0, posinf=0.0, neginf=0.0)
                    reducer = umap.UMAP(n_components=2, random_state=42,
                                        n_neighbors=15, min_dist=0.1,
                                        metric='euclidean')
                    reducer.fit(mat_sample_scaled)
                    # Batch transform in 1K-row batches
                    coords_list = []
                    for start in range(0, len(df), BATCH_SIZE):
                        end = min(start + BATCH_SIZE, len(df))
                        mat_batch = scaler.transform(data_matrix[start:end])
                        mat_batch = np.nan_to_num(mat_batch, nan=0.0, posinf=0.0, neginf=0.0)
                        coords_list.append(reducer.transform(mat_batch))
                    coords = np.vstack(coords_list)
                else:
                    data_scaled = StandardScaler().fit_transform(data_matrix)
                    # Handle NaN/Inf values (np is imported at module level)
                    data_scaled = np.nan_to_num(data_scaled, nan=0.0, posinf=0.0, neginf=0.0)
                    reducer = umap.UMAP(n_components=2, random_state=42,
                                        n_neighbors=15, min_dist=0.1,
                                        metric='euclidean')
                    coords = reducer.fit_transform(data_scaled)
                non_numeric = [c for c in df.columns if c not in numeric_cols and not c.startswith('_')]
                entity_col_umap = non_numeric[0] if non_numeric else df.columns[0]
                labels_list = cluster_entry.get('labels', [])
                # Save: create EntityProfile records with coordinates
                from analysis.models import EntityProfile as EP
                from analysis.models import ClusterResult as CR
                from analysis.models import AnalysisRun as AR
                from asgiref.sync import sync_to_async
                csv_glob = dataset_entry.get('metadata', {}).get('csv_glob', 'manual')
                run_obj = await sync_to_async(
                    AR.objects.filter(csv_glob=csv_glob).order_by('-id').first,
                    thread_sensitive=True
                )()
                if run_obj:
                    run_id_for_ep = run_obj.id
                    profiles = []
                    cr_cache = {}
                    for i, row in enumerate(df.iter_rows(named=True)):
                        if i >= len(coords): break
                        ev = str(row.get(entity_col_umap, ''))
                        if not ev: continue
                        lbl = labels_list[i] if i < len(labels_list) else -1
                        cache_key = f'{run_id_for_ep}_{lbl}'
                        if cache_key not in cr_cache:
                            cr_cache[cache_key] = CR.objects.filter(run_id=run_id_for_ep, cluster_label=lbl).first()
                        profiles.append(EP(
                            entity_value=ev, run_id=run_id_for_ep, cluster_label=lbl,
                            cluster=cr_cache[cache_key],
                            embedding_x=float(coords[i, 0]),
                            embedding_y=float(coords[i, 1]),
                        ))
                    if profiles:
                        EP.objects.bulk_create(profiles, ignore_conflicts=True, batch_size=1000)
            except Exception:
                import traceback, sys; print(f'[UMAP_ERR] {traceback.format_exc()}', file=sys.stderr)

        # Persist to Django ORM
        db_saved = False
        if save_to_db:
            try:
                from asgiref.sync import sync_to_async
                from analysis.models import AnalysisRun as AnalysisRunModel
                ce = cluster_entry  # local ref for closure
                async def _do_save(ce_inner):
                    if run_id is not None:
                        run = await sync_to_async(AnalysisRunModel.objects.get)(id=run_id)
                    else:
                        csv_glob_fallback = dataset_entry.get('metadata', {}).get('csv_glob', 'manual')
                        run = await sync_to_async(
                            AnalysisRunModel.objects.filter(csv_glob=csv_glob_fallback).order_by('-id').first,
                            thread_sensitive=True
                        )()
                        if run is None:
                            run = await sync_to_async(AnalysisRunModel.objects.create)(
                                csv_glob=csv_glob_fallback, status='running', total_flows=0,
                            )
                    run.status = 'completed'
                    run.cluster_count = ce_inner['n_clusters']
                    await sync_to_async(run.save)()
                    return await sync_to_async(_save_features_to_db)(cluster_result_id, all_features, run_id=run.id)
                db_saved = await _do_save(ce)
            except Exception as e:
                import traceback, sys; print(f'[DB_SAVE_ERROR] {e}\n{traceback.format_exc()}', file=sys.stderr)
                db_saved = False

    except Exception:
        raise

    result_id = _generate_id('feat')
    store.store_feature_result(
        result_id=result_id,
        features=all_features,
        parent_cluster_result_id=cluster_result_id,
    )

    return {
        'feature_result_id': result_id,
        'features': all_features,
        'n_features': len(all_features),
        'db_saved': db_saved,
    }


def _save_features_to_db(
    cluster_result_id: str,
    features: list[dict],
    run_id: int | None = None,
) -> bool:
    """Save extracted features to the Django ORM ClusterFeature model.

    If run_id is provided, use it directly (skip csv_glob lookup).
    """
    if run_id is not None:
        from analysis.models import AnalysisRun as AnalysisRun2
        run = AnalysisRun2.objects.filter(id=run_id).first()
        if run is None:
            return False
    else:
        # Try to find the ClusterResult record by run association
        store = SessionStore()
        cluster_entry = store.get_cluster_result(cluster_result_id)
        if cluster_entry is None:
            return False

        parent_ds_id = cluster_entry.get('parent_dataset_id', '')
        from analysis.models import AnalysisRun
        ds_entry = store.get_dataset(parent_ds_id)
        if ds_entry is None:
            return False

        csv_glob = ds_entry.get('metadata', {}).get('csv_glob', '')
        # Compute n_clusters from features before using it
        unique_global = sorted(set(f['cluster_label'] for f in features))
        n_clusters_global = len([l for l in unique_global if l >= 0])
        try:
            run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-created_at').first()
            if run is None:
                run = AnalysisRun.objects.create(
                    csv_glob=csv_glob,
                    status='completed',
                    cluster_count=n_clusters_global,
                )
        except Exception:
            return False

    from analysis.models import ClusterResult as ClusterResultModel

    # Compute actual cluster sizes from features (not from cluster_entry — avoids scope issues)
    label_counts = {}
    for feat in features:
        lbl = feat['cluster_label']
        label_counts[lbl] = label_counts.get(lbl, 0) + 1

    unique_labels = sorted(set(
        f['cluster_label'] for f in features
    ))
    # Compute n_clusters from features (excluding noise -1)
    n_clusters = len([l for l in unique_labels if l >= 0])
    for label in unique_labels:
        if label == -1:
            continue  # skip noise cluster
        label_features = [f for f in features if f['cluster_label'] == label]
        if not label_features:
            continue

        cluster_db, _ = ClusterResultModel.objects.get_or_create(
            run=run,
            cluster_label=label,
            defaults={'size': label_counts.get(label, 1)},
        )
        # Update size if already exists
        if label_counts.get(label) is not None:
            ClusterResultModel.objects.filter(
                run=run, cluster_label=label
            ).update(size=label_counts[label])

        from analysis.models import ClusterFeature as ClusterFeatureModel
        for feat in label_features:
            try:
                ClusterFeatureModel.objects.get_or_create(
                    cluster=cluster_db,
                    feature_name=feat['feature_name'],
                    defaults={
                        'mean': feat.get('mean'),
                        'std': feat.get('std'),
                        'median': None,
                        'p25': None,
                        'p75': None,
                        'missing_rate': None,
                        'distinguishing_score': feat.get('distinguishing_score'),
                        'distinguishing_method': feat.get('distinguishing_method', 'zscore'),
                    }
                )
            except Exception:
                continue  # skip problematic features

    # Update run counts
    run.cluster_count = n_clusters
    run.save(update_fields=['cluster_count'])
    return True

    return True


# ── 8. export_results ───────────────────────────────────────────────────

async def _handle_export_results(
    result_id: str,
    output_path: str,
    format: str = 'csv',
    overwrite: bool = False,
) -> dict:
    """Export a dataset (or result) to disk."""
    store = SessionStore()

    out = Path(output_path)
    file_paths: list[str] = []

    # Check if it's a dataset
    dataset_entry = store.get_dataset(result_id)
    if dataset_entry is not None:
        lf: pl.LazyFrame = dataset_entry['lazyframe']
        try:
            df = lf.collect(streaming=True)
            if out.suffix == '':
                out = out / f"export_{result_id}.{format}"

            if out.exists() and not overwrite:
                return {'error': f'Output exists and overwrite=False: {out}',
                        'truncated': False}

            out.parent.mkdir(parents=True, exist_ok=True)

            if format == 'csv':
                df.write_csv(out)
            elif format == 'parquet':
                df.write_parquet(out)
            elif format == 'json':
                df.write_json(out)
            else:
                return {'error': f'Unsupported format: {format}',
                        'truncated': False}

            file_paths.append(str(out.resolve()))
        except Exception:
            raise

        return {'file_paths': file_paths, 'format': format,
                'row_count': len(df)}

    # Check if it's a cluster result
    cluster_entry = store.get_cluster_result(result_id)
    if cluster_entry is not None:
        if out.suffix == '':
            out = out / f"cluster_{result_id}.json"
        out.parent.mkdir(parents=True, exist_ok=True)

        if out.exists() and not overwrite:
            return {'error': f'Output exists and overwrite=False: {out}',
                    'truncated': False}

        export_data = {
            'result_id': result_id,
            'n_clusters': cluster_entry['n_clusters'],
            'params': cluster_entry['params'],
            'labels_sample': cluster_entry['labels'][:1000],
        }
        with open(out, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, default=str)
        file_paths.append(str(out.resolve()))

        return {'file_paths': file_paths, 'format': 'json'}

    # Check feature result
    feat_entry = store.get_feature_result(result_id)
    if feat_entry is not None:
        if out.suffix == '':
            out = out / f"features_{result_id}.json"
        out.parent.mkdir(parents=True, exist_ok=True)

        if out.exists() and not overwrite:
            return {'error': f'Output exists and overwrite=False: {out}',
                    'truncated': False}

        with open(out, 'w', encoding='utf-8') as f:
            json.dump(feat_entry['features'], f, default=str)
        file_paths.append(str(out.resolve()))
        return {'file_paths': file_paths, 'format': 'json'}

    return {'error': f'No dataset or result found for ID: {result_id}',
            'truncated': False}


# ── 9. list_datasets ────────────────────────────────────────────────────

async def _handle_list_datasets(
    _filter_type: Optional[str] = None,
) -> dict:
    """List all active datasets and results."""
    store = SessionStore()
    datasets = store.list_datasets()
    return {'datasets': datasets, 'count': len(datasets)}


# ── 10. drop_dataset ────────────────────────────────────────────────────

async def _handle_drop_dataset(dataset_id: str) -> dict:
    """Remove a dataset/result and release memory."""
    store = SessionStore()

    dropped = (store.drop_dataset(dataset_id)
               or store.drop_cluster_result(dataset_id)
               or store.drop_feature_result(dataset_id))

    if dropped:
        return {'dropped': True, 'dataset_id': dataset_id}
    return {'error': f'No dataset or result found: {dataset_id}',
            'truncated': False}


# ── 11. clone_dataset ───────────────────────────────────────────────────

async def _handle_clone_dataset(dataset_id: str) -> dict:
    """Shallow-clone a dataset."""
    store = SessionStore()
    new_id = store.clone_dataset(dataset_id)
    if new_id is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    return {'original_dataset_id': dataset_id, 'cloned_dataset_id': new_id}


# ── 12. build_entity_profiles ────────────────────────────────────────────

_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)


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.

    Stores the result in the session store and returns the new dataset_id.
    """
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    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:
        # Auto-detect: find columns matching entity keywords
        matched = []
        for kw in _ENTITY_KEYWORDS:
            found = [c for c in col_names if kw.lower() in c.lower()]
            matched.extend(found)
        # Deduplicate while preserving order
        seen = set()
        group_cols = [c for c in matched if not (c in seen or seen.add(c))]
        if not group_cols:
            # Fallback: pick first non-numeric, non-internal column
            for c in col_names:
                if c.startswith('_'):
                    continue
                group_cols = [c]
                break
        # Limit to at most 3 entity columns
        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 ──────────────────────────────────
    # Always: count of flows
    agg_exprs: list[pl.Expr] = [
        pl.len().alias('flow_count'),
        pl.col(group_cols[0]).n_unique().alias(f'unique_{group_cols[0]}'),
    ]

    # Numeric aggregations (sum, mean)
    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}'))

    # Unique counts for string-like entity columns (dst_ip, dst_port, sni, etc.)
    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}'))

    # Also count distinct values of specific keyword-matching columns
    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 ───────────────────────────────────────────────
    try:
        entity_lf = lf.group_by(group_cols).agg(agg_exprs)
        # Estimate row count
        n_entities = entity_lf.select(pl.len()).collect(streaming=True).item()
    except Exception:
        raise

    # ── 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.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 anomaly scoring helpers (inlined from entity_aggregator) ─────

_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',
]


def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame:
    """Add anomaly scores using data-driven weights (SVD + IQR + percentile).

    Steps:
    1. Find which anomaly feature columns exist in the aggregated data
    2. Collect a sample (max 10K rows) for SVD weight learning
    3. Learn feature weights from SVD first component loadings
    4. Compute weighted proxy_score across ALL rows (vectorized, no per-row ops)
    5. Add percentile rank calibration
    6. Infer risk thresholds from data distribution
    """
    _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:
        # Step 1-2: Sample up to 10K rows for weight learning
        n_rows = lf.select(pl.len()).collect(streaming=True).item()
        sample_lf = lf if n_rows <= 10000 else lf.sample(n=10000, seed=42)
        df_sample = sample_lf.select(avail).collect(streaming=True)
        mat = df_sample.to_numpy()
        mat = np.nan_to_num(mat, nan=0.0)

        # Step 3: SVD weight learning
        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=42).fit(mat)
        raw_weights = np.abs(svd.components_[0])
        # Normalise to sum = 1
        weights = dict(zip(avail, raw_weights / raw_weights.sum()))

        # Step 4: Weighted sum expression (all vectorised, no loop)
        expr_parts = []
        for col, w in weights.items():
            if col == 'modern_tls_rate':
                # modern_tls_rate is good — invert for anomaly scoring
                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'))

        # Step 5: Percentile rank calibration (0-1, relative position)
        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()

        # Compute IQR thresholds from sample
        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

        # Step 6: Risk level labels (0.0-1.0)
        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)}
        import logging
        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, 10000), n_rows,
        )

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

    return lf


# ── 12. compute_scores ─────────────────────────────────────────────────────

async def _handle_compute_scores(dataset_id: str) -> dict:
    """Compute adaptive proxy/anomaly/threat scores on entity data."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf = entry['lazyframe']
    lf = _add_adaptive_scores(lf, [])
    store.store_dataset(dataset_id, lf, schema=entry.get('schema', {}),
                        metadata=entry.get('metadata', {}))
    return {'status': 'scored', 'dataset_id': dataset_id}


# ── 13. filter_and_cluster ─────────────────────────────────────────────────

async def _handle_filter_and_cluster(
    dataset_id: str,
    filters: list[dict],
    logic: str = 'and',
    algorithm: str = 'hdbscan',
    params: dict | None = None,
) -> dict:
    """Apply column-level filters then run clustering on the filtered result."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

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

    # ── Step 1: Apply filters ─────────────────────────────────────────────
    try:
        exprs = [_build_filter_expr(
            f['column'], f.get('op', 'eq'), f.get('value'),
        ) for f in filters]

        if exprs:
            if logic == 'and':
                combined = exprs[0]
                for e in exprs[1:]:
                    combined = combined & e
            else:
                combined = exprs[0]
                for e in exprs[1:]:
                    combined = combined | e
            filtered_lf = lf.filter(combined)
        else:
            filtered_lf = lf

        # Estimate row count (fast streaming count)
        try:
            row_count = filtered_lf.select(pl.lit(1).count()).collect(streaming=True).item()
        except Exception:
            row_count = None

    except Exception:
        raise

    # ── Step 2: Auto-detect numeric feature columns ───────────────────────
    try:
        available = filtered_lf.collect_schema()
        available_names = available.names()
        available_dtypes = available.dtypes()
        numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                          pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                          pl.Float32, pl.Float64)
        feature_cols = [
            available_names[i] for i, dt in enumerate(available_dtypes)
            if dt in numeric_dtypes and not available_names[i].startswith('_')
        ][:10]

        if not feature_cols:
            return {
                'error': 'No numeric columns found for clustering after filtering',
                'filtered_dataset_id': None,
                'truncated': False,
            }
    except Exception:
        raise

    # ── Step 3: Store filtered dataset ────────────────────────────────────
    filtered_dataset_id = _generate_id('flt')
    store.store_dataset(
        dataset_id=filtered_dataset_id,
        lazyframe=filtered_lf,
        schema=full_schema,
        metadata={
            'row_count': row_count,
            'parent_dataset_id': dataset_id,
            'filter_count': len(filters),
            'logic': logic,
        },
        parent_id=dataset_id,
    )

    # ── Step 4: Run clustering (inlined from _handle_run_clustering) ───────
    p = params or {}
    try:
        # Downsample before collect
        try:
            total_count = filtered_lf.select(pl.len()).collect(streaming=True).item()
        except Exception:
            total_count = 0

        clamp_lf = filtered_lf
        if total_count > 100000:
            clamp_lf = filtered_lf.collect(streaming=True).sample(n=100000, seed=42).lazy()

        df_features = clamp_lf.select(feature_cols).collect(streaming=True)

        # Drop all-NaN columns
        for col in df_features.columns:
            if df_features[col].is_null().all():
                df_features = df_features.drop(col)
                feature_cols = [c for c in feature_cols if c != col]

        if df_features.width == 0:
            return {'error': 'No valid feature columns after dropping all-NaN columns',
                    'filtered_dataset_id': filtered_dataset_id, 'truncated': False}

        data = df_features.to_numpy()

        # Fill remaining NaN with column mean
        if data.dtype.kind == 'f':
            col_mean = np.nanmean(data, axis=0)
            col_mean = np.nan_to_num(col_mean, nan=0.0)
            data = np.where(np.isnan(data), col_mean, data)

        # Too-few-samples guard
        if len(data) < 3:
            return {
                'cluster_result_id': _generate_id('clust'),
                'n_clusters': 0,
                'n_noise': 0,
                'quality_metrics': {'note': 'too few samples'},
                'filtered_dataset_id': filtered_dataset_id,
                'truncated': False,
            }

        # Low memory check
        low_memory = False
        try:
            import psutil
            mem = psutil.virtual_memory()
            low_memory = mem.available < 2 * 1024 ** 3
        except ImportError:
            pass

        if low_memory and len(data) > 100000:
            rng = np.random.default_rng(42)
            idx = rng.choice(len(data), 50000, replace=False)
            data = data[idx]

        # Downsample to 50K max
        _ds_idx = None
        if len(data) > 50000:
            rng = np.random.default_rng(42)
            _ds_idx = rng.choice(len(data), 50000, replace=False)
            data = data[_ds_idx]

        # Variance filtering
        if data.shape[1] > 1:
            variances = np.var(data, axis=0)
            keep_var = np.where(variances >= 1e-10)[0]
            if len(keep_var) == 0:
                keep_var = np.arange(data.shape[1])
            if len(keep_var) < data.shape[1]:
                data = data[:, keep_var]
                feature_cols = [feature_cols[i] for i in keep_var]

        # Standard scale
        from sklearn.preprocessing import StandardScaler
        data_scaled = StandardScaler().fit_transform(data)

        # Correlation filtering
        if data_scaled.shape[1] > 1:
            corr = np.corrcoef(data_scaled.T)
            upper = np.triu(np.abs(corr), 1)
            to_drop = [j for j in range(upper.shape[1]) if any(upper[:, j] > 0.95)]
            if to_drop:
                keep = [j for j in range(data_scaled.shape[1]) if j not in to_drop]
                data_scaled = data_scaled[:, keep]
                feature_cols = [feature_cols[j] for j in keep]

        # SVD dimensionality reduction (if >50 features)
        n_features = data_scaled.shape[1]
        if n_features > 50:
            from sklearn.decomposition import TruncatedSVD
            data_scaled = TruncatedSVD(n_components=50, random_state=42).fit_transform(data_scaled)

        # Hyperparameter auto-tuning
        n = len(data_scaled)
        if algorithm == 'hdbscan':
            min_cluster = p.get('min_cluster_size', max(3, min(50, n // 20)))
            p['min_cluster_size'] = min_cluster
            p['min_samples'] = p.get('min_samples', min_cluster)
            p['cluster_selection_epsilon'] = p.get('cluster_selection_epsilon', 0.0)
        else:
            p['min_cluster_size'] = min(
                p.get('min_cluster_size', 5),
                max(n // 2, 2),
            )

        # Fit model
        from sklearn.cluster import HDBSCAN, MiniBatchKMeans
        if algorithm == 'kmeans':
            n_clusters_param = p.get('n_clusters', 3)
            model = MiniBatchKMeans(
                n_clusters=n_clusters_param, random_state=42,
                batch_size=1024, n_init='auto',
            )
        else:
            model = HDBSCAN(
                min_cluster_size=p['min_cluster_size'],
                min_samples=p.get('min_samples', None),
                metric=p.get('metric', 'euclidean'),
                cluster_selection_epsilon=p.get('cluster_selection_epsilon', 0.0),
            )

        labels_array = model.fit_predict(data_scaled)
        n_clusters_found = int(len(set(labels_array)) - (1 if -1 in labels_array else 0))
        noise_count = int((labels_array == -1).sum()) if algorithm == 'hdbscan' else 0
        noise_ratio = round(noise_count / len(labels_array), 4) if algorithm == 'hdbscan' else 0.0

        # Quality metrics
        quality = {
            'n_clusters': n_clusters_found,
            'n_noise': noise_count,
            'noise_ratio': noise_ratio,
            'algorithm': algorithm,
        }

        if n_clusters_found > 1:
            try:
                from sklearn.metrics import silhouette_score
                score = silhouette_score(data_scaled, labels_array, random_state=42)
                quality['silhouette_score'] = round(float(score), 4)
            except Exception:
                quality['silhouette_score'] = None
            try:
                from sklearn.metrics import davies_bouldin_score
                dbs = davies_bouldin_score(data_scaled, labels_array)
                quality['davies_bouldin_score'] = round(float(dbs), 4)
            except Exception:
                quality['davies_bouldin_score'] = None
            try:
                from sklearn.metrics import calinski_harabasz_score
                chs = calinski_harabasz_score(data_scaled, labels_array)
                quality['calinski_harabasz_score'] = round(float(chs), 4)
            except Exception:
                quality['calinski_harabasz_score'] = None

        # Cluster sizes
        cluster_sizes = {}
        for label in np.unique(labels_array):
            cluster_sizes[int(label)] = int((labels_array == label).sum())
        quality['cluster_sizes'] = cluster_sizes

    except Exception:
        raise

    result_id = _generate_id('clust')
    store.store_cluster_result(
        result_id=result_id,
        labels=labels_array.tolist(),
        n_clusters=n_clusters_found,
        params={
            'algorithm': algorithm,
            'params': p,
            'random_state': 42,
            'feature_columns': feature_cols,
            '_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None,
        },
        parent_dataset_id=filtered_dataset_id,
    )

    return {
        'cluster_result_id': result_id,
        'filtered_dataset_id': filtered_dataset_id,
        'n_clusters': n_clusters_found,
        'n_noise': noise_count,
        'labels_sample': labels_array[:50].tolist(),
        'quality_metrics': quality,
        'feature_columns': feature_cols,
    }


# ── 14. detect_anomalies ──────────────────────────────────────────────────

async def _handle_detect_anomalies(dataset_id: str, contamination: float = 0.05) -> dict:
    """Run Isolation Forest on scored entity data."""
    from sklearn.ensemble import IsolationForest
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf: pl.LazyFrame = entry['lazyframe']
    schema = entry.get('schema', {})
    numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
    num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')]
    if not num_cols:
        return {'error': 'No numeric columns for anomaly detection', 'truncated': False}

    # Collect scored data (already sampled by pipeline head param)
    df = lf.select(num_cols).collect(streaming=True)
    mat = df.to_numpy()
    mat = np.nan_to_num(mat, nan=0.0)

    # Isolation Forest in chunks
    n = len(mat)
    chunk_size = min(50000, n)
    labels = np.zeros(n, dtype=int)
    for start in range(0, n, chunk_size):
        end = min(start + chunk_size, n)
        iso = IsolationForest(n_estimators=100, max_samples=min(256, chunk_size),
                              contamination=contamination, random_state=42, n_jobs=-1)
        labels[start:end] = iso.fit_predict(mat[start:end])

    anomaly_count = int((labels == -1).sum())
    return {
        'status': 'completed',
        'total_entities': n,
        'anomaly_count': anomaly_count,
        'anomaly_rate': round(anomaly_count / max(n, 1), 4),
        'dataset_id': dataset_id,
    }


# ── 15. visualize_anomalies ───────────────────────────────────────────────

async def _handle_visualize_anomalies(dataset_id: str) -> dict:
    """Generate UMAP-2D embedding for anomaly visualization."""
    from sklearn.preprocessing import StandardScaler
    import umap
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf: pl.LazyFrame = entry['lazyframe']
    schema = entry.get('schema', {})
    numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
    num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')]
    if len(num_cols) < 2:
        return {'error': 'Need at least 2 numeric columns for UMAP', 'truncated': False}

    # Sample for UMAP visualisation (max 10K)
    try:
        n_total = lf.select(pl.len()).collect(streaming=True).item()
    except Exception:
        n_total = 0
    sample_lf = lf if n_total <= 10000 else lf.collect(streaming=True).sample(n=10000, seed=42).lazy()
    df = sample_lf.select(num_cols).collect(streaming=True)
    mat = StandardScaler().fit_transform(df.to_numpy())
    mat = np.nan_to_num(mat, nan=0.0)
    reducer = umap.UMAP(n_components=2, random_state=42,
                        n_neighbors=15, min_dist=0.1, metric='euclidean')
    coords = reducer.fit_transform(mat)

    risk_col = next((c for c in schema if 'risk' in c.lower()), None)
    score_col = next((c for c in schema if 'proxy_score' in c.lower() or 'threat' in c.lower()), None)

    non_num = [c for c in df.columns if c not in num_cols and not c.startswith('_')]
    ent_col = non_num[0] if non_num else df.columns[0]
    entities = df.select(ent_col).to_series().to_list() if ent_col in df.columns else []

    scatter = [
        {'x': float(coords[i, 0]), 'y': float(coords[i, 1]),
         'entity': str(entities[i]) if i < len(entities) else '',
         'score': round(float(mat[i].mean()), 4)}
        for i in range(len(coords))
    ]
    return {
        'status': 'completed',
        'scatter_points': len(scatter),
        'scatter_data': scatter[:200],  # cap for LLM context
        'dataset_id': dataset_id,
    }


# ── 16. validate_data ──────────────────────────────────────────────────

async def _handle_validate_data(dataset_id: str, fix_issues: bool = False) -> dict:
    """Check data quality: schema consistency, missing values, type conflicts."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    issues = []
    try:
        df_sample = lf.head(1000).collect(streaming=True)
        for col in schema:
            null_count = df_sample[col].is_null().sum()
            if null_count > 500:
                issues.append(f'{col}: {null_count/10:.0f}% null')
            try:
                unique = df_sample[col].n_unique()
                if unique <= 1 and null_count == 0:
                    issues.append(f'{col}: constant value (no variation)')
            except Exception:
                pass  # n_unique not supported for this column type
    except Exception as exc:
        raise
    total_rows = lf.select(pl.len()).collect(streaming=True).item()
    return {
        'status': 'ok',
        'total_rows': total_rows,
        'columns': len(schema),
        'issues_found': len(issues),
        'issues': issues[:20],
        'recommendation': 'Try explore_distributions for detailed column analysis' if issues else 'No issues found',
    }


# ── 17. explore_distributions ──────────────────────────────────────────

async def _handle_explore_distributions(dataset_id: str, columns: list[str] = None,
                                        max_sample: int = 10000) -> dict:
    """Per-column distribution statistics."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}

    if not columns:
        columns = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types][:10]
    try:
        n = lf.select(pl.len()).collect(streaming=True).item()
        sample_lf = lf if n <= max_sample else lf.collect(streaming=True).sample(n=max_sample, seed=42).lazy()
        df = sample_lf.select(columns).collect(streaming=True)
    except Exception as exc:
        raise

    stats = {}
    for col in columns:
        try:
            s = df[col]
            stats[col] = {
                'dtype': str(s.dtype),
                'null_pct': round(float(s.is_null().mean() * 100), 2),
                'unique': int(s.n_unique()),
                'top_values': [str(v) for v in s.head(5).to_list()],
            }
            if s.dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64):
                stats[col]['min'] = float(s.min()) if s.min() is not None else None
                stats[col]['max'] = float(s.max()) if s.max() is not None else None
                stats[col]['mean'] = round(float(s.mean()), 4) if s.mean() is not None else None
                stats[col]['std'] = round(float(s.std()), 4) if s.std() is not None else None
        except Exception:
            stats[col] = {'error': 'failed to compute'}
    return {'columns': stats, 'total_rows': n, 'sampled': min(n, max_sample)}


# ── 18. find_outliers ─────────────────────────────────────────────────

async def _handle_find_outliers(dataset_id: str, columns: list[str] = None,
                                threshold: float = 3.0) -> dict:
    """Find statistical outliers using IQR method."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
    if not columns:
        columns = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types]
    try:
        n = lf.select(pl.len()).collect(streaming=True).item()
        sample_lf = lf if n <= 100000 else lf.collect(streaming=True).sample(n=100000, seed=42).lazy()
        df = sample_lf.select(columns).collect(streaming=True)
    except Exception as exc:
        raise

    outlier_counts = {}
    for col in columns:
        try:
            s = df[col]
            q1, q3 = s.quantile(0.25), s.quantile(0.75)
            iqr = q3 - q1
            lo = q1 - threshold * iqr
            hi = q3 + threshold * iqr
            n_out = ((s < lo) | (s > hi)).sum()
            if n_out > 0:
                outlier_counts[col] = int(n_out)
        except Exception:
            pass  # quantile not supported for this column type
    return {'outlier_counts': outlier_counts, 'total_sampled': len(df), 'threshold_iqr': threshold}


# ── 19. diagnose_clustering ───────────────────────────────────────────

async def _handle_diagnose_clustering(dataset_id: str, cluster_result_id: str) -> dict:
    """Diagnose clustering quality issues."""
    from sklearn.decomposition import TruncatedSVD
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    cluster_entry = store.get_cluster_result(cluster_result_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    if cluster_entry is None:
        return {'error': f'Cluster result not found: {cluster_result_id}', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
    num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')]

    n_clusters = cluster_entry.get('n_clusters', 0)
    n_noise = cluster_entry.get('n_noise', 0)
    quality = cluster_entry.get('quality_metrics', {})
    labels = cluster_entry.get('labels', [])
    params = cluster_entry.get('params', {})
    algo = params.get('algorithm', 'unknown')
    feature_cols = params.get('feature_columns', [])

    diagnosis = []
    if n_clusters <= 1:
        diagnosis.append('Only 1 cluster — data may lack structure or need feature scaling')
    if n_noise > 0.5 * len(labels) if labels else False:
        diagnosis.append(f'High noise ratio ({n_noise}/{len(labels)}) — try HDBSCAN with larger min_cluster_size')
    sil = quality.get('silhouette_score')
    if sil and sil < 0.2:
        diagnosis.append(f'Low silhouette ({sil:.3f}) — clusters overlap, try KMeans or feature selection')
    if len(feature_cols) < 3:
        diagnosis.append(f'Only {len(feature_cols)} features — add more columns for meaningful clusters')

    # SVD variance check
    try:
        df_sample = lf.select(num_cols).head(5000).collect(streaming=True)
        mat = df_sample.to_numpy()
        n_comp = min(mat.shape)
        svd = TruncatedSVD(n_components=n_comp, random_state=42).fit(mat)
        var_explained = svd.explained_variance_ratio_
        if var_explained[0] > 0.95:
            diagnosis.append(f'First SVD component explains {var_explained[0]:.0%} variance — data nearly 1D')
    except Exception:
        pass  # SVD fit failed — likely constant or all-null columns

    return {
        'diagnosis': diagnosis,
        'n_clusters': n_clusters,
        'n_noise': n_noise,
        'silhouette': sil,
        'algorithm': algo,
        'n_features': len(feature_cols),
        'suggestions': [
            f'Try algorithm={"kmeans" if algo == "hdbscan" else "hdbscan"}',
            'Increase min_cluster_size for fewer, larger clusters',
            'Use explore_distributions to check feature quality',
        ],
    }


# ── 20. compare_datasets ─────────────────────────────────────────────

async def _handle_compare_datasets(dataset_id_a: str, dataset_id_b: str) -> dict:
    """Compare two datasets."""
    store = SessionStore()
    entry_a, entry_b = store.get_dataset(dataset_id_a), store.get_dataset(dataset_id_b)
    if entry_a is None or entry_b is None:
        return {'error': 'One or both datasets not found', 'truncated': False}

    def _get_info(entry):
        schema = entry.get('schema', {})
        lf = entry['lazyframe']
        try:
            n = lf.select(pl.len()).collect(streaming=True).item()
        except Exception:
            n = -1
        return {'schema': set(schema.keys()), 'row_count': n, 'columns': len(schema)}

    info_a, info_b = _get_info(entry_a), _get_info(entry_b)
    common_cols = info_a['schema'] & info_b['schema']
    only_a = info_a['schema'] - info_b['schema']
    only_b = info_b['schema'] - info_a['schema']
    return {
        'dataset_a': {'id': dataset_id_a, **info_a},
        'dataset_b': {'id': dataset_id_b, **info_b},
        'common_columns': len(common_cols),
        'columns_only_in_a': list(only_a)[:10] if only_a else [],
        'columns_only_in_b': list(only_b)[:10] if only_b else [],
        'row_diff': info_a['row_count'] - info_b['row_count'],
        'identical_schema': len(only_a) == 0 and len(only_b) == 0,
    }


# ── 21. export_debug_sample ───────────────────────────────────────────

async def _handle_export_debug_sample(dataset_id: str, max_rows: int = 50) -> dict:
    """Export raw data sample as JSON for debugging."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    try:
        df = entry['lazyframe'].head(max_rows).collect(streaming=True)
        rows = df.to_dict(as_series=False)
        # Convert to list-of-dicts, serialize safely
        result = []
        for i in range(len(df)):
            row = {}
            for col in df.columns:
                val = rows[col][i]
                if isinstance(val, (float, int, str, bool)):
                    row[col] = val
                elif val is None:
                    row[col] = None
                else:
                    row[col] = str(val)
            result.append(row)
        return {'rows': result, 'count': len(result), 'columns': df.columns}
    except Exception as exc:
        raise


# ── 22. repair_schema ─────────────────────────────────────────────────

async def _handle_repair_schema(dataset_id: str) -> dict:
    """Attempt to repair schema mismatches by normalising column names."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    # Normalise: replace hyphens/underscores/dots, lowercase
    col_map = {}
    for c in schema:
        norm = c.lower().replace('-', '_').replace('.', '_')
        col_map[c] = norm
    if len(set(col_map.values())) < len(col_map):
        # Conflicts exist — can't auto-repair
        return {'status': 'conflict', 'message': 'Normalised names conflict', 'column_map': col_map}
    # Rename columns
    renames = {old: new for old, new in col_map.items() if old != new}
    if renames:
        lf = lf.rename(renames)
        new_schema = {col_map.get(c, c): dtype for c, dtype in schema.items()}
        store.store_dataset(dataset_id, lf, schema=new_schema, metadata=entry.get('metadata', {}))
        return {'status': 'repaired', 'renamed': len(renames), 'column_map': col_map}
    return {'status': 'no_change', 'message': 'Column names already normalised'}


# ── 23. analyze_patterns ──────────────────────────────────────────────

async def _handle_analyze_patterns(dataset_id: str, top_n: int = 10) -> dict:
    """Read-only: traffic pattern analysis."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': 'Dataset not found', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    n = lf.select(pl.len()).collect(streaming=True).item()
    df = lf.head(100000).collect(streaming=True)

    result = {'total_rows': n, 'sampled': min(n, 100000)}

    # Top source IPs
    for col_name, label in [(':ips', 'src_ips'), (':ipd', 'dst_ips'),
                              ('snam', 'snis'), ('cnam', 'cert_names')]:
        if col_name in df.columns:
            try:
                top = df[col_name].value_counts(sort=True).head(top_n)
                result[label] = {
                    'top': [{'value': str(r[col_name]), 'count': int(r['count'])}
                            for r in top.iter_rows(named=True)],
                    'unique': int(df[col_name].n_unique()),
                }
            except Exception:
                pass  # value_counts not supported for column, skip

    # Port distribution
    for port_col in [':prs', ':prd']:
        if port_col in df.columns:
            try:
                top = df[port_col].value_counts(sort=True).head(top_n)
                result[f'{port_col}_dist'] = {
                    'top': [{'port': int(r[port_col]), 'count': int(r['count'])}
                            for r in top.iter_rows(named=True)],
                    'unique': int(df[port_col].n_unique()),
                }
            except Exception:
                pass  # port distribution extraction failed — skip

    # TLS version breakdown
    for tls_col in ['0ver']:
        if tls_col in df.columns:
            try:
                vc = df[tls_col].value_counts(sort=True)
                result['tls_versions'] = [{'version': str(r[tls_col]), 'count': int(r['count'])}
                                          for r in vc.iter_rows(named=True)]
            except Exception:
                pass  # TLS version value_counts not supported — skip

    return result


# ── 24. analyze_temporal ─────────────────────────────────────────────

async def _handle_analyze_temporal(dataset_id: str, timestamp_column: str = None) -> dict:
    """Read-only: temporal traffic pattern analysis."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': 'Dataset not found', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})

    if not timestamp_column:
        for cand in ['timestamp', 'time', '8dbd', '8did']:
            if cand in schema:
                timestamp_column = cand
                break
    if not timestamp_column:
        return {'error': 'No timestamp column found', 'truncated': False}

    n = lf.select(pl.len()).collect(streaming=True).item()
    return {'total_rows': n, 'timestamp_column': timestamp_column,
            'note': 'Temporal analysis requires parsed timestamps; raw values shown in explore_distributions'}


# ── 24b. analyze_fft ─────────────────────────────────────────────────

async def _handle_analyze_fft(
    dataset_id: str,
    timestamp_column: str,
    top_n: int = 5,
) -> dict:
    """FFT spectrum analysis on a timestamp column to detect periodic patterns.

    Computes the Fast Fourier Transform of sorted timestamp-derived values,
    returning dominant frequencies, spectral centroid, and a periodicity score.
    """
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

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

    # ── Resolve timestamp column ──────────────────────────────────────
    col = None
    # Try exact match first
    if timestamp_column in schema:
        col = timestamp_column
    else:
        # Try case-insensitive match
        col_lower = timestamp_column.lower()
        for c in schema:
            if c.lower() == col_lower:
                col = c
                break
    if col is None:
        return {'error': f"Timestamp column '{timestamp_column}' not found in dataset. Available: {list(schema.keys())[:20]}",
                'truncated': False}

    # ── Extract and sort values ───────────────────────────────────────
    df = lf.select(pl.col(col).alias('__ts__')).collect(streaming=True)

    n_raw = len(df)
    if n_raw < 4:
        return {'error': f'Need at least 4 samples for FFT; got {n_raw}', 'truncated': False}

    # Determine dtype and convert to numeric values
    ts_series = df['__ts__']
    dtype = ts_series.dtype
    str_dtypes = (pl.Utf8, pl.String)

    if dtype in str_dtypes:
        # Try parsing as datetime string, then convert to epoch seconds
        try:
            ts_parsed = ts_series.str.strptime(pl.Datetime, format=None, strict=False)
            ts_numeric = ts_parsed.cast(pl.Int64) / 1_000_000_000.0  # ns → seconds
        except Exception:
            # Fallback: try to cast to float directly
            try:
                ts_numeric = ts_series.cast(pl.Float64)
            except Exception as e:
                return {'error': f'Cannot convert string column to numeric or datetime: {e}',
                        'truncated': False}
    elif dtype in (pl.Datetime, pl.Date):
        ts_numeric = ts_series.cast(pl.Int64) / 1_000_000_000.0
    else:
        # Float, Int types — cast to float
        try:
            ts_numeric = ts_series.cast(pl.Float64)
        except Exception as e:
            return {'error': f'Cannot cast column to float: {e}', 'truncated': False}

    # Build DataFrame with values, drop nulls, sort
    df_sorted = pl.DataFrame({'__val__': ts_numeric}).drop_nulls().sort('__val__')
    n = len(df_sorted)
    if n < 4:
        return {'error': f'After removing nulls, only {n} samples remain (need ≥4)', 'truncated': False}

    values = df_sorted['__val__'].to_numpy().astype(np.float64)

    # ── FFT computation ────────────────────────────────────────────────
    # Center by subtracting mean (DC offset removal)
    values_centered = values - np.mean(values)
    fft_complex = np.fft.fft(values_centered)
    fft_magnitudes = np.abs(fft_complex)
    freqs = np.fft.fftfreq(n)

    # Take positive half (excluding DC at index 0)
    pos_mask = freqs > 0
    freqs_pos = freqs[pos_mask]
    mags_pos = fft_magnitudes[pos_mask]

    if len(freqs_pos) == 0:
        return {'error': 'Insufficient positive frequencies for FFT analysis', 'truncated': False}

    # ── Find top-N peaks ──────────────────────────────────────────────
    peak_indices = np.argsort(mags_pos)[::-1][:min(top_n, len(mags_pos))]
    top_periods = []
    for idx in peak_indices:
        freq = float(freqs_pos[idx])
        mag = float(mags_pos[idx])
        # Convert frequency to period (in sample-index units)
        period = float(1.0 / freq) if freq > 1e-10 else float('inf')
        # Map to time-domain period if data is temporal
        if n > 1:
            time_span = float(values[-1] - values[0])
            sample_period = time_span / max(n - 1, 1)
            time_period = period * sample_period
        else:
            time_period = float('inf')
        top_periods.append({
            'period': round(period, 2),
            'magnitude': round(mag, 4),
            'frequency': round(freq, 6),
            'time_domain_period_seconds': round(time_period, 2) if time_period != float('inf') else None,
        })

    # ── Spectral centroid (weighted mean frequency) ────────────────────
    mag_sum = float(np.sum(mags_pos))
    if mag_sum > 0:
        spectral_centroid = float(np.sum(freqs_pos * mags_pos) / mag_sum)
    else:
        spectral_centroid = 0.0

    # ── Periodicity score: ratio of top peak to mean magnitude ────────
    mean_mag = float(np.mean(mags_pos)) if len(mags_pos) > 0 else 1.0
    max_mag = float(mags_pos[peak_indices[0]]) if len(peak_indices) > 0 else 0.0
    periodicity_score = round(max_mag / max(mean_mag, 1e-10), 4)

    # ── Return decimated data for LLM context window ──────────────────
    # Decimate to at most 200 frequency-magnitude pairs
    decimate_step = max(1, len(freqs_pos) // 200)
    freqs_decimated = freqs_pos[::decimate_step].tolist()
    mags_decimated = mags_pos[::decimate_step].tolist()

    return {
        'frequencies': [round(float(f), 6) for f in freqs_decimated],
        'magnitudes': [round(float(m), 4) for m in mags_decimated],
        'top_periods': top_periods,
        'spectral_centroid': round(spectral_centroid, 6),
        'periodicity_score': periodicity_score,
        'n_samples': n,
        'n_raw': n_raw,
        'nan_dropped': n_raw - n,
        'timestamp_column': col,
    }


# ── 25. analyze_tls_health ────────────────────────────────────────────

async def _handle_analyze_tls_health(dataset_id: str) -> dict:
    """Read-only: TLS security assessment."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': 'Dataset not found', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    n = lf.select(pl.len()).collect(streaming=True).item()
    df = lf.head(100000).collect(streaming=True)

    result = {'total_rows': n, 'sampled': min(n, 100000)}

    if '0ver' in df.columns:
        vc = df['0ver'].value_counts(sort=True)
        versions = {}
        for r in vc.iter_rows(named=True):
            v = str(r['0ver']).replace(' ', '')
            if v in ('0303', '0304'):
                versions['tls_1_2_plus'] = versions.get('tls_1_2_plus', 0) + int(r['count'])
            else:
                versions['tls_legacy'] = versions.get('tls_legacy', 0) + int(r['count'])
        total = sum(versions.values())
        result.update({
            'tls_1_2_plus_pct': round(versions.get('tls_1_2_plus', 0) / max(total, 1) * 100, 1),
            'tls_legacy_pct': round(versions.get('tls_legacy', 0) / max(total, 1) * 100, 1),
            'tls_legacy_versions': [{'version': str(r['0ver']), 'count': int(r['count'])}
                                     for r in vc.iter_rows(named=True)
                                     if str(r['0ver']).replace(' ', '') not in ('0303', '0304')],
        })

    if 'snam' in df.columns:
        empty_sni = df['snam'].is_null().sum() + df['snam'].cast(pl.Utf8).is_in(['', 'null', 'None']).sum()
        result['sni_missing_pct'] = round(float(empty_sni) / max(len(df), 1) * 100, 1)

    if 'cipher-suite' in df.columns:
        weak_ciphers = ['rc4', 'cbc', 'null', 'export']
        total_ciphers = len(df)
        weak_count = 0
        for w in weak_ciphers:
            try:
                weak_count += df['cipher-suite'].str.to_lowercase().str.contains(w).sum()
            except Exception:
                pass  # str.contains not supported for cipher-suite column — skip
        result['weak_cipher_pct'] = round(float(weak_count) / max(total_ciphers, 1) * 100, 1) if total_ciphers else 0

    if result.get('tls_legacy_pct', 0) > 20:
        result['warning'] = 'High proportion of legacy TLS — possible downgrade attacks'
    if result.get('sni_missing_pct', 0) > 50:
        result['warning'] = 'Most flows lack SNI — possible IP-direct connections (proxy/VPN)'
    return result


# ── 26. analyze_geo_distribution ─────────────────────────────────────

async def _handle_analyze_geo_distribution(dataset_id: str) -> dict:
    """Read-only: geographic traffic analysis."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': 'Dataset not found', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    n = lf.select(pl.len()).collect(streaming=True).item()
    df = lf.head(100000).collect(streaming=True)

    result = {'total_rows': n, 'sampled': min(n, 100000)}

    for geo_col, label in [('scnt', 'src_countries'), ('dcnt', 'dst_countries')]:
        if geo_col in df.columns:
            try:
                vc = df[geo_col].value_counts(sort=True).head(15)
                result[label] = [{'country': str(r[geo_col]), 'count': int(r['count'])}
                                 for r in vc.iter_rows(named=True)]
                result[f'{label}_unique'] = int(df[geo_col].n_unique())
            except Exception:
                pass  # country distribution not computable — skip

    if 'scnt' in df.columns and 'dcnt' in df.columns:
        try:
            same = (df['scnt'] == df['dcnt']).sum()
            result['domestic_pct'] = round(float(same) / max(len(df), 1) * 100, 1)
            result['international_pct'] = round(100.0 - result['domestic_pct'], 1)
        except Exception:
                pass  # domestic/international comparison failed — skip

    return result


# ── 27. analyze_entity_detail ────────────────────────────────────────

async def _handle_analyze_entity_detail(dataset_id: str, entity_value: str) -> dict:
    """Read-only: deep-dive into a single entity."""
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': 'Dataset not found', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})

    # Find entity column (first non-numeric, non-internal column)
    numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
    entity_col = None
    for c in schema:
        if c.startswith('_') or schema[c].split('(')[0].strip() in numeric_types or '.' in c:
            continue
        entity_col = c
        break
    if not entity_col:
        entity_col = list(schema.keys())[0]

    try:
        filtered = lf.filter(pl.col(entity_col).cast(pl.Utf8).str.contains(entity_value, literal=True))
        n = filtered.select(pl.len()).collect(streaming=True).item()
        if n == 0:
            return {'error': f"Entity '{entity_value}' not found in column '{entity_col}'", 'truncated': False}
        df = filtered.collect(streaming=True)
    except Exception as exc:
        raise

    result = {
        'entity_value': entity_value,
        'entity_column': entity_col,
        'flow_count': n,
        'columns': list(df.columns[:15]),
    }

    # Numeric summary
    for c in df.columns:
        if df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64):
            try:
                result[f'{c}_stats'] = {
                    'mean': round(float(df[c].mean()), 2),
                    'min': float(df[c].min()) if df[c].min() is not None else None,
                    'max': float(df[c].max()) if df[c].max() is not None else None,
                }
            except Exception:
                pass  # numeric stats not computable for this column — skip

    return result


# ── 28. compute_distance_matrix ──────────────────────────────────────────

def _indent_function_body(body: str, indent: str = '    ') -> str:
    """Indent each non-empty line of a code block for wrapping in a def."""
    return '\n'.join(
        indent + line if line.strip() else line
        for line in body.split('\n')
    )


async def _handle_compute_distance_matrix(
    dataset_id: str,
    python_function: str,
    columns: Optional[list[str]] = None,
    max_sample: int = 10000,
) -> dict:
    """Evaluate a user-provided Python function on each row of a dataset.

    Wraps the function body as ``def distance_fn(row): <body>``, executes
    it per-row in a restricted namespace (``math``, ``numpy`` only), and
    returns a distance-score summary plus a sample of scored rows.
    """
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    lf: pl.LazyFrame = entry['lazyframe']

    # Collect data (streaming for large datasets)
    df = lf.collect(streaming=True)

    # Determine columns to pass to the function
    if columns is None:
        numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                          pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                          pl.Float32, pl.Float64)
        columns = [c for c in df.columns
                   if df[c].dtype in numeric_dtypes]

    # Filter to columns that actually exist
    available = set(df.columns)
    columns = [c for c in columns if c in available]
    if not columns:
        return {'error': 'No valid columns found in dataset',
                'truncated': False}

    n_total = len(df)

    # Sample if larger than max_sample
    if n_total > max_sample:
        df_sampled = df.sample(n=max_sample, seed=42)
        n_evaluated = max_sample
    else:
        df_sampled = df
        n_evaluated = n_total

    # Build list-of-dicts for the selected columns
    rows = df_sampled.select(columns).to_dicts()

    # ── Safe execution ────────────────────────────────────────────────────
    safe_builtins = {
        'abs': abs, 'len': len, 'range': range,
        'min': min, 'max': max, 'sum': sum, 'round': round,
        'int': int, 'float': float, 'str': str, 'bool': bool,
        'list': list, 'dict': dict, 'tuple': tuple,
        'isinstance': isinstance, 'type': type,
        'True': True, 'False': False, 'None': None,
        'enumerate': enumerate, 'zip': zip, 'map': map,
        'filter': filter, 'sorted': sorted, 'reversed': reversed,
    }

    safe_ns = {
        '__builtins__': safe_builtins,
        'math': __import__('math'),
        'np': __import__('numpy'),
    }

    # Wrap the user's function body in a def statement
    body = python_function.strip()
    if body.startswith('def '):
        # User provided the full function definition
        full_code = body
    else:
        indented = _indent_function_body(body)
        full_code = f'def distance_fn(row):\n{indented}'

    try:
        exec(full_code, safe_ns)
        # Resolve: the function name might be 'distance_fn' or whatever the user named it
        distance_fn = safe_ns.get('distance_fn')
        if distance_fn is None:
            # Fallback: try to find any callable in safe_ns
            for _name, _obj in safe_ns.items():
                if callable(_obj) and _name not in ('math', 'np', '__builtins__'):
                    distance_fn = _obj
                    break
        if distance_fn is None:
            return {'error': 'Could not locate the compiled function',
                    'truncated': False}
    except SyntaxError as e:
        return {'error': f'Function syntax error: {e}', 'truncated': False}
    except Exception as e:
        return {'error': f'Function compilation failed: {e}',
                'truncated': False}

    # ── Apply function per row ────────────────────────────────────────────
    scores: list[float | None] = []
    errors = 0
    for row in rows:
        try:
            score = distance_fn(row)
            if isinstance(score, (int, float)):
                scores.append(float(score))
            else:
                scores.append(None)
                errors += 1
        except Exception:
            scores.append(None)
            errors += 1

    valid_scores = [s for s in scores if s is not None]
    if not valid_scores:
        return {'error': 'No valid scores computed — all rows raised errors',
                'truncated': False}

    # ── Distance summary ──────────────────────────────────────────────────
    arr = np.array(valid_scores)
    summary = {
        'min': float(np.min(arr)),
        'max': float(np.max(arr)),
        'mean': float(np.mean(arr)),
        'std': float(np.std(arr)),
        'count': len(arr),
        'errors': errors,
        'total_rows_in_sample': len(rows),
        'matrix_shape': [len(rows), len(columns)],
    }

    # Sample of scored rows (first 30, keep score + first 3 values)
    scored_sample = []
    for i, row in enumerate(rows[:30]):
        if i < len(scores):
            scored_sample.append({
                'index': i,
                'values': {c: row.get(c) for c in columns[:3]},
                'score': scores[i],
            })

    return {
        'summary': summary,
        'scored_sample': scored_sample,
        'columns_used': columns,
        'total_rows_in_dataset': n_total,
        'rows_evaluated': n_evaluated,
    }
