from __future__ import annotations

"""online-mining-v2 可迁移脚本的共享逻辑。

脚本只做只读 ELK 查询和轻量格式转换，输出统一 JSON，方便被不同 Agent 调用。
"""

import argparse
import importlib.util
import json
import os
import re
import sys
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any


class OnlineMiningError(ValueError):
    """线上挖掘参数或数据错误。"""


SCRIPT_DIR = Path(__file__).resolve().parent
SKILL_DIR = SCRIPT_DIR.parent
REPO_ROOT = SKILL_DIR.parent.parent
ELK_FETCH_DIR = SKILL_DIR.parent / "elk-fetch"


SOURCE_CONFIGS: dict[str, dict[str, Any]] = {
    "main": {
        "profile": "default",
        "index": "arch-flat-nlp-log-f-*",
        "id_field": "request_id",
        "time_field": "timestamp",
    },
    "pre_processing": {
        "profile": "default",
        "index": "pre-processing*",
        "id_field": "requestId",
        "time_field": "timestamp",
    },
}


def load_json_payload(argv: list[str] | None = None) -> dict[str, Any]:
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument("--input", "-i", help="JSON 参数文件；不传则从 stdin 读取")
    args = parser.parse_args(argv)
    text = Path(args.input).expanduser().read_text(encoding="utf-8") if args.input else sys.stdin.read()
    if not text.strip():
        raise OnlineMiningError("input JSON is required")
    payload = json.loads(text)
    if not isinstance(payload, dict):
        raise OnlineMiningError("input JSON must be an object")
    return payload


def emit_success(payload: dict[str, Any]) -> None:
    print(json.dumps({"ok": True, **payload}, ensure_ascii=False, separators=(",", ":")))


def emit_error(exc: BaseException) -> None:
    print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False, separators=(",", ":")))


def load_elk_query_module():
    path = ELK_FETCH_DIR / "elk_query.py"
    if not path.exists():
        raise OnlineMiningError(f"elk-fetch script not found: {path}")
    spec = importlib.util.spec_from_file_location("elk_query", path)
    if spec is None or spec.loader is None:
        raise OnlineMiningError(f"cannot load elk_query.py from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def build_client(source: str):
    elk = load_elk_query_module()
    config = source_config(source)
    profiles = elk.load_profiles()
    return elk.build_client(profiles, str(config["profile"]))


def source_config(source: str) -> dict[str, Any]:
    if source not in SOURCE_CONFIGS:
        raise OnlineMiningError(f"unknown source: {source}")
    return SOURCE_CONFIGS[source]


def date_range_ms(date_str: str | None) -> tuple[int, int]:
    if date_str:
        date_obj = datetime.strptime(str(date_str), "%Y%m%d")
        start_ms = int(date_obj.timestamp() * 1000)
        end_ms = int((date_obj + timedelta(days=1)).timestamp() * 1000) - 1
    else:
        end = datetime.now()
        start_ms = int((end - timedelta(hours=48)).timestamp() * 1000)
        end_ms = int(end.timestamp() * 1000)
    return start_ms, end_ms


def deep_parse(value: Any) -> Any:
    if isinstance(value, str):
        stripped = value.strip()
        if (stripped.startswith("{") and stripped.endswith("}")) or (
            stripped.startswith("[") and stripped.endswith("]")
        ):
            try:
                return deep_parse(json.loads(stripped))
            except Exception:
                return value
        return value
    if isinstance(value, dict):
        return {key: deep_parse(item) for key, item in value.items()}
    if isinstance(value, list):
        return [deep_parse(item) for item in value]
    return value


def get_path(obj: Any, dotted: str, default: Any = None) -> Any:
    cur = obj
    for part in dotted.split("."):
        if not part:
            continue
        if isinstance(cur, dict):
            cur = cur.get(part)
        elif isinstance(cur, list):
            try:
                cur = cur[int(part)]
            except (ValueError, IndexError):
                return default
        else:
            return default
    return default if cur is None else cur


def compact_text(value: Any, limit: int = 500) -> str:
    if value is None:
        return ""
    text = json.dumps(value, ensure_ascii=False) if isinstance(value, (dict, list)) else str(value)
    text = re.sub(r"\s+", " ", text).strip()
    return text if len(text) <= limit else text[: limit - 3] + "..."


def as_list(value: Any) -> list[Any]:
    if value is None:
        return []
    if isinstance(value, list):
        return value
    return [value]


def string_values(value: Any) -> list[str]:
    return [str(item) for item in as_list(value) if str(item).strip()]


def match_text_filters(text: str, contains: Any = None, regex: str | None = None) -> bool:
    for item in string_values(contains):
        if item not in text:
            return False
    if regex and not re.search(regex, text):
        return False
    return True


def unique_strings(items: list[Any]) -> list[str]:
    seen: set[str] = set()
    result: list[str] = []
    for item in items:
        text = str(item).strip()
        if not text or text in seen:
            continue
        seen.add(text)
        result.append(text)
    return result


def extract_case(source: str, doc: dict[str, Any]) -> dict[str, Any]:
    parsed = deep_parse(doc)
    if source == "main":
        return extract_main_case(parsed)
    if source == "pre_processing":
        return extract_pre_processing_case(parsed)
    raise OnlineMiningError(f"unknown source: {source}")


def extract_main_case(doc: dict[str, Any]) -> dict[str, Any]:
    device = doc.get("device")
    return {
        "request_id": doc.get("request_id"),
        "timestamp": doc.get("timestamp"),
        "query": doc.get("query"),
        "session_id": doc.get("session_id"),
        "device_id": doc.get("device_id"),
        "device": device,
        "domain": doc.get("domain"),
        "func": doc.get("func") or doc.get("func_name"),
        "text": doc.get("to_speak") or doc.get("text") or doc.get("display_text"),
        "intention": doc.get("intention"),
        "llm_agent_info": get_path(doc, "intention.intent_arbitrator_info.llm_agent_info"),
        "score_domains": get_path(doc, "intention.intent_arbitrator_info.score_domains"),
        "hit_rules": get_path(doc, "intention.intent_arbitrator_info.hit_rules"),
        "raw": doc,
    }


def extract_pre_processing_case(doc: dict[str, Any]) -> dict[str, Any]:
    body = doc.get("responseBoby") or {}
    core = get_path(body, "nodes.0.core", {})
    dispatch = get_path(core, "dispatch_large_model_result", {}) or {}
    dispatch_input = get_path(dispatch, "dispatchLargeModelInput", {}) or {}
    planning_debug = get_path(dispatch, "planningDebugInfo", {}) or {}
    properties0 = get_path(body, "nodes.0.properties.0", {}) or {}
    debug = properties0.get("debug") if isinstance(properties0, dict) else {}
    excellent_domains = get_path(core, "excellent_domains_result", []) or []
    domain_infos = get_path(debug, "domain_infos", []) if isinstance(debug, dict) else []
    return {
        "request_id": doc.get("requestId"),
        "timestamp": doc.get("timestamp"),
        "query": get_path(core, "query.query"),
        "planning_result": planning_debug.get("planningOriginalResult") if isinstance(planning_debug, dict) else None,
        "code": dispatch.get("code") if isinstance(dispatch, dict) else None,
        "prompt_model": dispatch_input.get("promptModel") if isinstance(dispatch_input, dict) else None,
        "hit_rules": dispatch_input.get("hitRuleList") if isinstance(dispatch_input, dict) else None,
        "excellent_domains": excellent_domains,
        "candidate_domains": unique_strings(
            [get_path(item, "domain") for item in as_list(excellent_domains)]
            + [get_path(item, "domain") for item in as_list(domain_infos)]
        ),
        "domain_infos": domain_infos,
        "agent_llm_intent": properties0.get("agentLLMIntent") if isinstance(properties0, dict) else None,
        "prompt_string": get_path(debug, "promptString") if isinstance(debug, dict) else None,
        "planning_prompt": get_path(debug, "planningPrompt") if isinstance(debug, dict) else None,
        "raw": doc,
    }


def case_matches(case: dict[str, Any], filters: dict[str, Any]) -> bool:
    query = str(case.get("query") or "")
    if not match_text_filters(query, filters.get("query_contains"), filters.get("query_regex")):
        return False
    for key in ("domain", "func", "planning_result", "code"):
        expected = string_values(filters.get(key))
        if expected and str(case.get(key) or "") not in expected:
            return False
        contains = filters.get(f"{key}_contains")
        if contains and not match_text_filters(str(case.get(key) or ""), contains):
            return False
    if filters.get("prompt_contains") and not match_text_filters(
        str(case.get("prompt_model") or case.get("prompt_string") or case.get("planning_prompt") or ""),
        filters.get("prompt_contains"),
    ):
        return False
    candidate_domains = set(string_values(case.get("candidate_domains")))
    expected_candidates = set(string_values(filters.get("candidate_domain")))
    if expected_candidates and not expected_candidates.intersection(candidate_domains):
        return False
    return True


def search_docs(
    *,
    source: str,
    date: str | None,
    size: int,
    query_filters: list[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
    config = source_config(source)
    client = build_client(source)
    start_ms, end_ms = date_range_ms(date)
    filters = [{"range": {str(config["time_field"]): {"gte": start_ms, "lte": end_ms}}}]
    filters.extend(deepcopy(query_filters or []))
    body = {
        "query": {"bool": {"filter": filters}},
        "size": size,
        "sort": [{str(config["time_field"]): {"order": "desc"}}],
    }
    response = client.search(index=str(config["index"]), body=body)
    return [deep_parse(hit.get("_source", {})) for hit in response.get("hits", {}).get("hits", [])]


def fetch_one_by_request_id(source: str, request_id: str, date: str | None = None) -> dict[str, Any] | None:
    config = source_config(source)
    id_field = str(config["id_field"])
    # 精确查优先用 .keyword；部分索引字段本身就是 keyword，所以再兜底原字段。
    attempts = [
        {"term": {f"{id_field}.keyword": request_id}},
        {"term": {id_field: request_id}},
        {"wildcard": {f"{id_field}.keyword": {"value": f"{request_id}*"}}},
        {"wildcard": {id_field: {"value": f"{request_id}*"}}},
    ]
    for query_filter in attempts:
        docs = search_docs(source=source, date=date, size=1, query_filters=[query_filter])
        if docs:
            return docs[0]
    return None


def fetch_session_turns(
    *,
    session_id: str,
    before_timestamp: int,
    date: str | None = None,
    limit: int = 10,
) -> list[dict[str, Any]]:
    """按主表 session_id 拉当前请求之前的多轮上下文。"""

    if not session_id or not before_timestamp:
        return []
    docs: list[dict[str, Any]] = []
    # arch-flat-nlp-log-f-* 里 session_id 在部分索引没有 .keyword 子字段，
    # 所以必须先试 keyword，再兜底原字段。
    for session_field in ("session_id.keyword", "session_id"):
        filters = [
            {"term": {session_field: session_id}},
            {"range": {"timestamp": {"lt": before_timestamp}}},
        ]
        docs = search_docs(source="main", date=date, size=max(limit * 3, limit), query_filters=filters)
        if docs:
            break
    cases = [extract_main_case(doc) for doc in docs]
    cases = [
        case
        for case in cases
        if str(case.get("session_id") or "") == session_id
        and isinstance(case.get("timestamp"), (int, str))
        and int(case.get("timestamp") or 0) < before_timestamp
        and str(case.get("query") or "").strip()
    ]
    cases.sort(key=lambda item: int(item.get("timestamp") or 0))
    return cases[-limit:]


def fetch_session_turns_for_case(case: dict[str, Any], *, date: str | None = None, limit: int = 10) -> list[dict[str, Any]]:
    session_id = str(case.get("session_id") or "").strip()
    try:
        timestamp = int(case.get("timestamp") or 0)
    except (TypeError, ValueError):
        timestamp = 0
    return fetch_session_turns(session_id=session_id, before_timestamp=timestamp, date=date, limit=limit)


def fetch_by_terms(source: str, *, field: str, value: str, date: str | None = None, size: int = 10) -> list[dict[str, Any]]:
    """小范围精确字段查询，供脚本探测 session/request 字段时使用。"""

    filters = [{"term": {field: value}}]
    return search_docs(source=source, date=date, size=size, query_filters=filters)


def fetch_one_by_request_id_legacy(source: str, request_id: str, date: str | None = None) -> dict[str, Any] | None:
    """保留旧函数名兼容外部脚本；新代码请用 fetch_one_by_request_id。"""

    return fetch_one_by_request_id(source, request_id, date)


def write_optional_jsonl(path: str | None, rows: list[dict[str, Any]]) -> str | None:
    if not path:
        return None
    output = resolve_portable_path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
    return str(output)


def resolve_portable_path(path: str) -> Path:
    """解析 online-mining-v2 脚本路径，并优先映射到当前会话目录。

    在 ZK Data Agent 中，python_exec 会注入 PYTHON_EXEC_SCRATCHPAD。
    用户和 skill 文档里写的 scratchpad/、output/、input/ 都应该落在
    当前会话下，不能误写到项目根目录。
    """

    raw = Path(path).expanduser()
    if raw.is_absolute():
        return raw
    scratchpad = os.environ.get("PYTHON_EXEC_SCRATCHPAD")
    if scratchpad:
        session_root = Path(scratchpad).expanduser().parent
        parts = raw.parts
        if parts:
            head, *tail = parts
            tail_path = Path(*tail) if tail else Path()
            if head in {"scratchpad", "scratch"}:
                return (Path(scratchpad).expanduser() / tail_path).resolve()
            if head in {"output", "outputs"}:
                return (session_root / "output" / tail_path).resolve()
            if head in {"input", "inputs"}:
                return (session_root / "input" / tail_path).resolve()
    return (REPO_ROOT / raw).resolve()