"""LLM orchestrator — intelligent tool chaining for TLS data analysis.
Designed for 32B models (DeepSeek/Qwen/Llama).  Short prompts,
aggressive truncation for tool results, step-by-step guidance."""

import json, logging, time, urllib.request, urllib.error, asyncio
from typing import Optional
from analysis.tool_registry import handle_call

logger = logging.getLogger('tianxuan.llm')

# ── Tool descriptions (shared with MCP server) ─────────────────────────

TOOLS = [
    # ── Core pipeline ──
    {"type": "function", "function": {"name": "profile_data",
        "description": "检查数据集列：类型、空值比例、唯一值计数、数值统计。首先调用以了解数据，再决定后续步骤。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "build_entity_profiles",
        "description": "按实体列（源IP、SNI等）分组原始流。生成每个实体的特征（流数量、TLS版本、端口）。需要已概要分析的数据。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "entity_column": {"type": "string"}, "entity_columns": {"type": "array", "items": {"type": "string"}}, "auto_detect": {"type": "boolean"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "compute_scores",
        "description": "使用自适应权重学习在实体数据上计算自适应代理/异常/威胁评分。在build_entity_profiles之后调用。添加proxy_score/risk_level/threat_score列。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "run_clustering",
        "description": "将实体数据聚类为行为组。需要实体数据集ID。探索性分析选择hdbscan（自动检测聚类数）；固定簇数选择kmeans。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "cluster_columns": {"type": "array", "items": {"type": "string"}}, "algorithm": {"type": "string", "enum": ["hdbscan", "kmeans"]}}, "required": ["dataset_id", "cluster_columns"]}}},
    {"type": "function", "function": {"name": "extract_features",
        "description": "识别每个聚类的区分特征（z-score或anova）。在run_clustering之后调用。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "cluster_result_id": {"type": "string"}, "method": {"type": "string", "enum": ["zscore", "anova"]}}, "required": ["dataset_id", "cluster_result_id"]}}},
    {"type": "function", "function": {"name": "detect_anomalies",
        "description": "运行孤立森林（Isolation Forest）检测异常实体。在compute_scores之后调用。返回异常数量和比率。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "contamination": {"type": "number", "description": "Expected anomaly rate 0.01-0.1 (default 0.05)"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "visualize_anomalies",
        "description": "生成UMAP-2D散点数据用于前端展示。按风险等级着色。在compute_scores或detect_anomalies之后调用。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
    # ── Diagnostic (call when pipeline fails or results look wrong) ──
    {"type": "function", "function": {"name": "validate_data",
        "description": "[诊断] 检查数据质量：空值比例、常量列、schema问题。当加载失败或结果异常时调用。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "fix_issues": {"type": "boolean"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "explore_distributions",
        "description": "[诊断] 每列分布统计：唯一值计数、空值百分比、最小/最大/均值/标准差、top-10值。当聚类失败或需要了解列值时调用。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "columns": {"type": "array", "items": {"type": "string"}}, "max_sample": {"type": "integer"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "find_outliers",
        "description": "[诊断] 使用IQR方法查找极值。当explore_distributions显示可疑极端值时调用。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "columns": {"type": "array", "items": {"type": "string"}}, "threshold": {"type": "number"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "diagnose_clustering",
        "description": "[诊断] 当聚类结果不正常（全部噪声、只有1个聚类、轮廓系数低）时：诊断原因并建议修复方案。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "cluster_result_id": {"type": "string"}}, "required": ["dataset_id", "cluster_result_id"]}}},
    {"type": "function", "function": {"name": "compare_datasets",
        "description": "[诊断] 比较两个数据集，检查处理步骤是否损坏了数据。",
        "parameters": {"type": "object", "properties": {"dataset_id_a": {"type": "string"}, "dataset_id_b": {"type": "string"}}, "required": ["dataset_id_a", "dataset_id_b"]}}},
    {"type": "function", "function": {"name": "export_debug_sample",
        "description": "[诊断] 导出前N行JSON用于手动检查。当无法理解工具错误时调用。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "max_rows": {"type": "integer"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "repair_schema",
        "description": "[诊断] 通过规范化列名修复schema不匹配。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
    # Read-only analysis tools (deep-dive, no side effects)
    {"type": "function", "function": {"name": "analyze_patterns",
        "description": "[分析] 流量模式摘要——top IP、端口、TLS分布、协议混合。只读。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "top_n": {"type": "integer"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "analyze_tls_health",
        "description": "[分析] TLS安全评估——旧版比例、弱加密套件、SNI缺失率。只读。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "analyze_geo_distribution",
        "description": "[分析] 流量地理分布——top国家、国内vs国际、不可能行程。只读。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "analyze_entity_detail",
        "description": "[分析] 深入分析单个实体——所有流、TLS、端口、目的地、异常评分。只读。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "entity_value": {"type": "string"}}, "required": ["dataset_id", "entity_value"]}}},
    {"type": "function", "function": {"name": "analyze_temporal",
        "description": "[分析] 随时间变化的流量模式——小时/日流量、高峰时段。只读。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "timestamp_column": {"type": "string"}}, "required": ["dataset_id"]}}},
    {"type": "function", "function": {"name": "analyze_fft",
        "description": "[分析] FFT频谱分析——检测时间戳列的隐藏周期模式（每日/每周等）。返回频率、幅度、top周期、频谱质心、周期性评分。只读。在analyze_temporal之后用于检测周期性。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "timestamp_column": {"type": "string"}, "top_n": {"type": "integer", "description": "Top peaks to return (default 5)"}}, "required": ["dataset_id", "timestamp_column"]}}},
    # Custom LLM-powered computation
    {"type": "function", "function": {"name": "compute_distance_matrix",
        "description": "[LLM-POWERED] 你编写一个Python函数体的body（`def distance_fn(row: dict) -> float:`），系统逐行执行并返回距离分数统计。用于现有工具不能表达的度量任务。参数：python_function（函数体）、columns（列名列表）、max_sample（最大采样行数，默认10000）。",
        "parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "python_function": {"type": "string"}, "columns": {"type": "array", "items": {"type": "string"}}, "max_sample": {"type": "integer"}}, "required": ["dataset_id", "python_function"]}}},
]


SYSTEM_PROMPT = """你是TLS流量分析AI。你的任务是用可用工具分析网络流数据。
初始消息包含各列的元数据（TlsDB类型、中文含义、空值率、唯一值数、最常见值），可直接参考这些信息决策。未知列的类型/含义标注为"未知"。

## 策略
1. 首先调用`profile_data`了解数据集（列、类型、规模）。
2. 根据概要分析结果决定下一步：
   - 有实体列？→ `build_entity_profiles`
   - 需要异常评分？→ `compute_scores`（实体画像之后）
   - 需要聚类？→ `run_clustering` + `extract_features`
3. 如果工具失败或返回异常结果→调用[诊断]工具排查。
4. 最后用中文总结分析发现。

## 规则
- 始终检查工具结果后再决定下一步。
- 如果聚类返回0个聚类或全部噪声→调用`diagnose_clustering`。
- 如果实体检测失败→在候选列上调用`explore_distributions`。
- 绝不用完全相同参数重复调用同一工具。
- 绝不编造参数——使用工具结果中的实际值。
- 保持回复简洁。总结用中文，工具调用用英文。
- 每次分析会话最多15次工具调用。

## 可用工具
核心（修改数据）：profile_data, build_entity_profiles, compute_scores, run_clustering, extract_features, detect_anomalies, visualize_anomalies
诊断（出错时使用）：validate_data, explore_distributions, find_outliers, diagnose_clustering, compare_datasets, export_debug_sample, repair_schema
分析（只读，随时安全调用）：analyze_patterns, analyze_tls_health, analyze_geo_distribution, analyze_entity_detail, analyze_temporal, analyze_fft
LLM驱动（自定义计算）：compute_distance_matrix — 分析你的数据列，编写一个计算每行评分的distance_fn函数体，然后调用此工具执行并返回汇总统计。

## LLM驱动的compute_distance_matrix工作流
- 分析列→编写Python函数体→调用compute_distance_matrix→基于评分进行聚类或过滤
- 函数签名格式为`def distance_fn(row: dict) -> float:`。你只需编写函数体（例如`return abs(row.get('col1', 0)) + row.get('col2', 0)`）
- 可用命名空间：math、numpy（作为np）。禁止危险模块（os、sys、subprocess等）。
- 当现有工具无法表达所需度量时使用。

## 何时调用分析工具
- 决定实体列前→analyze_patterns（查看top IP）
- 聚类之前→analyze_tls_health（检查TLS质量）
- proxy_score高时→analyze_entity_detail调查该实体
- 怀疑地理异常时→analyze_geo_distribution
- 安全评估→analyze_tls_health
- 这些是只读工具——随时可调用，无副作用"""


class LLMConfig:
    def __init__(self, base_url="", api_key="", model="deepseek-v4-flash",
                 max_tokens=4096, temperature=0.1):
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        self.model = model
        self.max_tokens = max_tokens
        self.temperature = temperature


def call_llm(messages, config):
    """Call the LLM API with tool definitions. Returns parsed JSON response."""
    payload = json.dumps({
        "model": config.model,
        "messages": messages,
        "tools": TOOLS,
        "tool_choice": "auto",
        "max_tokens": config.max_tokens,
        "temperature": config.temperature,
    }).encode()
    headers = {"Content-Type": "application/json"}
    if config.api_key:
        headers["Authorization"] = f"Bearer {config.api_key}"
    req = urllib.request.Request(
        f"{config.base_url}/chat/completions",
        data=payload, headers=headers, method='POST',
    )
    logger.info(f'[LLM] model={config.model}')
    start = time.monotonic()
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            d = json.loads(resp.read().decode())
            elapsed = int((time.monotonic() - start) * 1000)
            logger.info(f'[LLM] latency_ms={elapsed}')
            return d
    except urllib.error.HTTPError as e:
        logger.error(f'[LLM] HTTP {e.code}: {e.read().decode()[:200]}')
        return None
    except Exception as e:
        logger.error(f'[LLM] call failed: {e}')
        return None


def _make_dataset_context(dataset_id: str) -> str:
    """Build rich column-level context for LLM: types from TlsDB, meanings,
    null rate, unique count, and most frequent value (mode).

    Graceful fallback when tls_ref_data table is empty or stats computation
    fails — column metadata degrades to bare names with '未知' labels."""
    try:
        import polars as pl
        from analysis.session_store import SessionStore
        from analysis.tls_ref import get_tls_ref_dict

        store = SessionStore()
        entry = store.get_dataset(dataset_id)
        if entry is None:
            return f"Dataset '{dataset_id}' not found."

        schema: dict = entry.get('schema', {})
        metadata: dict = entry.get('metadata', {})
        row_count = metadata.get('row_count', 0)
        col_names = list(schema.keys())

        # TlsDB reference lookup (graceful fallback to 未知)
        try:
            tls_ref = get_tls_ref_dict()
        except Exception:
            tls_ref = {}

        lf = entry.get('lazyframe')

        lines = []
        row_info = f"{row_count} rows" if row_count else "unknown rows"
        lines.append(f"Dataset '{dataset_id}': {row_info}, {len(col_names)} columns.")

        if lf is None or not col_names:
            return '\n'.join(lines)

        # Limit displayed columns to avoid overwhelming LLM context
        display_cols = col_names[:50]

        try:
            # Resolve actual row count if metadata is missing
            n_rows = row_count
            if not n_rows:
                n_rows = lf.select(pl.len()).collect().item()

            # ── Batch 1: null_count + n_unique in one collect ──────────
            stat_exprs: list[pl.Expr] = []
            stat_map: dict[str, tuple[int, str]] = {}  # idx -> (null_alias, uniq_alias)
            for i, c in enumerate(display_cols):
                null_a = f"_sn_{i}"
                uniq_a = f"_su_{i}"
                stat_exprs.append(pl.col(c).null_count().alias(null_a))
                stat_exprs.append(pl.col(c).n_unique().alias(uniq_a))
                stat_map[c] = (null_a, uniq_a)

            stats_df = lf.select(stat_exprs).collect()

            # ── Batch 2: sample data for mode computation ──────────────
            SAMPLE_CAP = 30000
            if n_rows > SAMPLE_CAP:
                sample_df = lf.slice(0, SAMPLE_CAP).select(display_cols).collect()
            else:
                sample_df = lf.select(display_cols).collect()

            # ── Build per-column line ─────────────────────────────────
            lines.append("Columns:")
            for i, c in enumerate(display_cols):
                ref = tls_ref.get(c, {})
                data_type = ref.get('data_type', '未知')
                meaning = ref.get('meaning', '未知')

                null_a, uniq_a = stat_map[c]
                null_cnt = stats_df.item(0, null_a) if null_a in stats_df else 0
                null_pct = f"{(null_cnt / n_rows * 100):.1f}%" if n_rows > 0 else "0%"
                uniq = stats_df.item(0, uniq_a) if uniq_a in stats_df else 0

                # Mode (most frequent value)
                try:
                    mode_series = sample_df.select(pl.col(c).mode().first()).get_column(c)
                    if len(mode_series) > 0 and mode_series[0] is not None:
                        mode_str = str(mode_series[0])[:30]
                    else:
                        mode_str = "N/A"
                except Exception:
                    mode_str = "N/A"

                lines.append(
                    f"  {c} ({data_type}, {meaning}): "
                    f"null={null_pct}, unique={uniq}, mode={mode_str}"
                )
        except Exception:
            # Stats failed — fall back to bare column list with tls_ref labels
            lines.append("Columns (stats unavailable):")
            for c in display_cols:
                ref = tls_ref.get(c, {})
                data_type = ref.get('data_type', '未知')
                meaning = ref.get('meaning', '未知')
                lines.append(f"  {c} ({data_type}, {meaning})")

        if len(col_names) > 50:
            lines.append(f"  ... and {len(col_names) - 50} more columns")

        return '\n'.join(lines)
    except Exception:
        return f"Dataset '{dataset_id}'"


def run_llm_pipeline(dataset_id: str, entity_col: str = "",
                     config: Optional[LLMConfig] = None,
                     callback=None) -> dict:
    """Run an LLM-driven analysis pipeline.

    The LLM autonomously chooses which tools to call and in what order.
    The system prompt provides strategy guidance; the LLM decides the rest.

    Args:
        dataset_id: Source dataset from load_data.
        entity_col: Optional pre-detected entity column.
        config: LLM configuration (model, API key, etc.).
        callback: Optional progress callback(step, tool_name, result).

    Returns:
        dict with 'result' (final summary) and 'steps' (tool calls made).
    """
    if config is None:
        from config import get_config
        cfg = get_config()
        config = LLMConfig(
            cfg.llm.base_url, cfg.llm.api_key, cfg.llm.model,
        )

    # Build initial context with rich column metadata
    ctx = _make_dataset_context(dataset_id)
    user_msg = ctx + "\n\n"
    if entity_col:
        user_msg += f"The entity column is '{entity_col}'. "
    else:
        user_msg += "Detect the entity column automatically. "
    user_msg += "Start by profiling, then decide next steps based on the data."

    msgs = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_msg},
    ]

    # Track tool call history for smarter guidance
    called_tools: set[str] = set()
    consecutive_errors = 0

    for step in range(25):
        resp = call_llm(msgs, config)
        if resp is None:
            if callback:
                callback(step, "thinking", "(LLM unavailable)")
            consecutive_errors += 1
            if consecutive_errors >= 3:
                return {"error": "LLM unavailable after 3 retries", "steps": step}
            continue
        consecutive_errors = 0

        msg = resp["choices"][0]["message"]
        content = msg.get("content", "")

        # Pass thinking text to callback
        if callback:
            callback(step, "thinking", content)

        # No tool calls → LLM is done
        if "tool_calls" not in msg or not msg["tool_calls"]:
            if callback:
                callback(step, "complete", content)
            return {"result": content or "Analysis complete", "steps": step + 1}

        msgs.append({"role": "assistant", "content": content, "tool_calls": msg["tool_calls"]})

        for tc in msg["tool_calls"]:
            fn = tc.get("function", {})
            name = fn.get("name", "")
            raw_args = fn.get("arguments", "{}")
            try:
                args = json.loads(raw_args)
            except json.JSONDecodeError:
                args = {}
                logger.warning(f'[LLM] bad JSON args from tool {name}: {raw_args[:100]}')

            # Check for repeated tool call with same params
            tool_key = f"{name}({json.dumps(args, sort_keys=True)})"
            if tool_key in called_tools and name in ('profile_data', 'explore_distributions'):
                logger.warning(f'[LLM] repeated tool call: {tool_key}')
                msgs.append({
                    "role": "user",
                    "content": f"Tool '{name}' was already called with these arguments. Try a different tool or different parameters.",
                })
                continue
            called_tools.add(tool_key)

            logger.info(f'[TOOL] step={step} name={name} args={json.dumps(args)[:200]}')
            if callback:
                callback(step, name, None, meta={"args": args})

            try:
                result = asyncio.run(handle_call(name, args))
            except Exception as e:
                raise SystemExit(f"MCP tool error: {e}") from e

            if callback:
                callback(step, name, result, meta={"args": args})

            # Truncate result for LLM context window
            result_str = json.dumps(result, default=str, ensure_ascii=False)
            if len(result_str) > 2000:
                # Keep structure but truncate large arrays
                if isinstance(result, dict):
                    for k in list(result.keys()):
                        v = result[k]
                        if isinstance(v, (list, dict)) and len(json.dumps(v)) > 500:
                            result[k] = f"<{len(v)} items, truncated>"
                elif isinstance(result, list) and len(result) > 20:
                    result = result[:10] + [f"<+{len(result)-10} more>"]
                result_str = json.dumps(result, default=str, ensure_ascii=False)[:2000]

            msgs.append({
                "role": "tool",
                "tool_call_id": tc.get("id", ""),
                "content": result_str,
            })

            # If tool returned an error, guide the LLM to retry or diagnose
            if isinstance(result, dict) and 'error' in result:
                logger.warning(f'[TOOL] {name} error: {result["error"]}')
                diag_map = {
                    'build_entity_profiles': 'explore_distributions',
                    'run_clustering': 'diagnose_clustering',
                    'load_data': 'repair_schema',
                }
                suggested = diag_map.get(name, 'validate_data')
                msgs.append({
                    "role": "user",
                    "content": f"Tool '{name}' failed: {result['error'][:200]}. Call '{suggested}' to diagnose, then retry with corrected parameters.",
                })

    if callback:
        callback(25, "max_steps", None)
    return {"error": "Reached max 25 tool calls", "steps": 25}
