"""
ELK / Elasticsearch 日志查询工具。

按 request id 查 ELK。支持：
- 多账号 profile（不同业务用不同 ES 账号）
- 多业务 preset（业务语义 → 索引/字段/过滤条件的映射）
- 深度 JSON 字符串自动解析（rejectInfo / message 等场景）
- 按点号路径提取嵌套字段（--json-path）

账号密码从同目录 profiles.json 加载（团队共用账号已 commit 进仓库；可用 ELK_FETCH_PROFILES 指向别处覆盖）。
业务→索引→路径的详细说明见同目录 INDEX_CATALOG.md。

用法：
    python elk_query.py <request_id>
        [--preset main|reject|kwfree|micar]
        [--date YYYYMMDD]
        [--index PATTERN] [--field FIELD] [--time-field NAME] [--profile KEY]
        [--size N]
        [--fields f1,f2,...]
        [--json-path a.b.c]        # 提取嵌套字段，自动解 JSON 字符串
        [--raw]                    # 关闭自动 JSON 字符串解析
"""
import argparse
import json
import os
import sys
import warnings
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Optional

import urllib3
from elasticsearch import Elasticsearch
from elasticsearch.exceptions import (
    ConnectionError as ESConnectionError,
    ElasticsearchWarning,
    NotFoundError,
    RequestError,
    AuthorizationException,
    AuthenticationException,
)

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
warnings.filterwarnings("ignore", category=ElasticsearchWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

ES_PORT = int(os.getenv("ES_PORT", "80"))
ES_HOST_FALLBACK = os.getenv("ES_HOST", "")

SCRIPT_DIR = Path(__file__).resolve().parent
PROFILES_PATH = Path(os.getenv("ELK_FETCH_PROFILES", SCRIPT_DIR / "profiles.json"))


def load_profiles() -> Dict[str, Dict[str, str]]:
    """从 profiles.json 读取账号配置。缺文件 / 格式错误时给出可操作的错误信息。"""
    if not PROFILES_PATH.exists():
        raise FileNotFoundError(
            f"profiles.json not found at {PROFILES_PATH}. "
            f"Override path via ELK_FETCH_PROFILES=/path/to/profiles.json"
        )
    try:
        with PROFILES_PATH.open("r", encoding="utf-8") as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f"profiles.json is not valid JSON: {e}")
    if not isinstance(data, dict) or not data:
        raise ValueError("profiles.json must be a non-empty JSON object keyed by profile name")
    for key, prof in data.items():
        if not isinstance(prof, dict) or "user" not in prof or "password" not in prof:
            raise ValueError(f"profile '{key}' missing required fields 'user' / 'password'")
    return data


# 业务→索引/字段/额外过滤 的映射。time_field 为毫秒时间戳字段名。
# 每条 preset 对应一个明确的业务语义，详见 INDEX_CATALOG.md。
PRESETS: Dict[str, Dict[str, Any]] = {
    "main": {
        "profile": "default",
        "index": "arch-flat-nlp-log-f-*",
        "field": "request_id",
        "time_field": "timestamp",
        "extra_filters": [],
    },
    "reject": {
        # 后置拒识模块的日志，内含调用免唤醒判决（KEY_WORD_FREE_RESTRIC）的结果
        "profile": "reject",
        "index": "aiservice_duplex_rejection_lcs_log*",
        "field": "requestId",
        "time_field": "time",
        "extra_filters": [],
    },
    "kwfree": {
        # 后处理阶段调用免唤醒模块自身的日志。message 是 JSON 字符串
        "profile": "default",
        "index": "nlp_post_processing_lcs-*",
        "field": "requestId",
        "time_field": "timestamp",
        "extra_filters": [{"match_phrase": {"moduleName": "keyword-free-log"}}],
    },
    "micar": {
        # 小米汽车小爱 OneTrack 埋点（端到端可用性 / 性能埋点等 tip）。
        # 索引按日 rolling 且非常大（~300GB/天），强烈建议传 --date 收敛日期。
        # key_value 是序列化 JSON 字符串，脚本会自动深度解析。
        "profile": "micar",
        "index": "onetrack_xiaoai_micar*",
        "field": "request_id",
        "time_field": "timestamp",
        "extra_filters": [],
    },
}


def build_client(profiles: Dict[str, Dict[str, str]], profile_key: str) -> Elasticsearch:
    if profile_key not in profiles:
        available = ", ".join(sorted(profiles.keys())) or "(none)"
        raise KeyError(f"profile '{profile_key}' not in profiles.json (available: {available})")
    prof = profiles[profile_key]
    host = prof.get("host") or ES_HOST_FALLBACK
    if not host:
        raise ValueError(
            f"profile '{profile_key}' has no 'host' and ES_HOST env var is not set"
        )
    return Elasticsearch(
        hosts=[f"http://{host}:{ES_PORT}"],
        http_auth=(prof["user"], prof["password"]),
        timeout=60,
        max_retries=3,
    )


def build_query(
    request_id: str,
    field: str,
    time_field: str,
    date_str: Optional[str],
    size: int,
    extra_filters: List[Dict[str, Any]],
) -> Dict[str, Any]:
    if date_str:
        date_obj = datetime.strptime(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_ms = int(datetime.now().timestamp() * 1000)
        start_ms = int((datetime.now() - timedelta(hours=48)).timestamp() * 1000)

    filters: List[Dict[str, Any]] = [
        {"wildcard": {field: {"value": f"{request_id}*"}}},
    ]
    filters.extend(deepcopy(extra_filters))

    return {
        "query": {
            "bool": {
                "must": [{"range": {time_field: {"gte": start_ms, "lte": end_ms}}}],
                "filter": filters,
            }
        },
        "size": size,
        "sort": [{time_field: {"order": "desc"}}],
    }


def deep_parse(v: Any) -> Any:
    """递归把看起来是 JSON 的字符串展开成 dict/list。"""
    if isinstance(v, str):
        s = v.strip()
        if (s.startswith("{") and s.endswith("}")) or (s.startswith("[") and s.endswith("]")):
            try:
                return deep_parse(json.loads(s))
            except Exception:
                return v
        return v
    if isinstance(v, dict):
        return {k: deep_parse(val) for k, val in v.items()}
    if isinstance(v, list):
        return [deep_parse(x) for x in v]
    return v


def get_path(obj: Any, dotted: str) -> 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 None
        else:
            return None
    return cur


def project(
    source: Dict[str, Any],
    fields: Optional[List[str]],
    json_path: Optional[str],
) -> Any:
    if json_path:
        return {json_path: get_path(source, json_path)}
    if fields:
        return {k: get_path(source, k) for k in fields}
    return source


def run(args: argparse.Namespace, profiles: Dict[str, Dict[str, str]]) -> Dict[str, Any]:
    preset = PRESETS[args.preset]
    profile_key = args.profile or preset["profile"]
    index_pattern = args.index or preset["index"]
    field_name = args.field or preset["field"]
    time_field = args.time_field or preset["time_field"]
    extra_filters = preset.get("extra_filters", [])

    body = build_query(
        args.request_id, field_name, time_field, args.date, args.size, extra_filters
    )
    try:
        client = build_client(profiles, profile_key)
    except (KeyError, ValueError) as e:
        return {"error": "ProfileError", "message": str(e)}

    try:
        resp = client.search(index=index_pattern, body=body)
    except AuthenticationException as e:
        return {"error": "AuthenticationException", "profile": profile_key, "message": str(e)}
    except AuthorizationException as e:
        return {"error": "AuthorizationException", "profile": profile_key, "index": index_pattern, "message": str(e)}
    except NotFoundError as e:
        return {"error": "IndexNotFound", "index": index_pattern, "message": str(e)}
    except RequestError as e:
        return {"error": "QueryError", "message": str(e)}
    except ESConnectionError as e:
        return {"error": "ConnectionError", "message": str(e)}

    hits_raw = resp.get("hits", {}).get("hits", [])
    fields = [f.strip() for f in args.fields.split(",")] if args.fields else None

    records = []
    for h in hits_raw:
        src = h.get("_source", {})
        if not args.raw:
            src = deep_parse(src)
        records.append(project(src, fields, args.json_path))

    return {
        "request_id": args.request_id,
        "preset": args.preset,
        "profile": profile_key,
        "index": index_pattern,
        "field": field_name,
        "time_field": time_field,
        "date": args.date or "past-48h",
        "total": len(records),
        "hits": records,
    }


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Query ELK by request id (multi-profile, deep-JSON aware)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="业务→索引映射详见同目录 INDEX_CATALOG.md",
    )
    p.add_argument("request_id", help="目标 request id（支持前缀通配）")
    p.add_argument("--preset", choices=list(PRESETS.keys()), default="main",
                   help="预置方案：" + " | ".join(PRESETS.keys()))
    p.add_argument("--profile", help="覆盖账号 profile（key 来自 profiles.json）")
    p.add_argument("--index", help="覆盖索引 pattern")
    p.add_argument("--field", help="覆盖 request id 字段名")
    p.add_argument("--time-field", help="覆盖时间字段名")
    p.add_argument("--date", help="指定日期 YYYYMMDD，不填默认近 48 小时")
    p.add_argument("--size", type=int, default=20, help="最多返回记录数（默认 20）")
    p.add_argument("--fields", help="只输出指定字段，逗号分隔，支持 a.b.c 嵌套")
    p.add_argument("--json-path", help="只提取某条嵌套路径的值")
    p.add_argument("--raw", action="store_true", help="关闭自动 JSON 字符串解析")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    try:
        profiles = load_profiles()
    except (FileNotFoundError, ValueError) as e:
        print(json.dumps({"error": "ConfigError", "message": str(e)}, indent=2, ensure_ascii=False))
        return 2
    result = run(args, profiles)
    print(json.dumps(result, indent=2, ensure_ascii=False))
    return 0 if "error" not in result else 1


if __name__ == "__main__":
    sys.exit(main())
