"""
搜索功能的Function Call包装器
"""

import json
import os
from datetime import datetime
import re
from typing import Optional, Union, List, Dict
from ketacli.sdk.ai.function_call import function_registry
from ketacli.sdk.ai.client import AIClient
from ketacli.sdk.base.search import search_spl, search_spl_meta, search_summary
from ketacli.sdk.output.output import search_result_output, list_output, get_asset_output, make_table
from ketacli.sdk.output.format import format_table
from ketacli.sdk.request.list import list_assets_request
from ketacli.sdk.base.client import request_get
from rich.console import Console
import logging

console = Console(markup=False)
logger = logging.getLogger("ketacli.textual.ai_helpers")

def validate_spl_value_quotes(spl: str) -> None:
    """通用校验：SPL 中所有出现“字段值”的场景必须使用双引号。

    原则：
    - 字段名用单引号（例如：'host'）。
    - 字段值用双引号（例如："10-0-1"，包含双引号时可用三双引号）。
    - 校验范围覆盖：
      * search/search2 主段中的条件（如 repo、start、end 及任意 key=value）
      * where 子句中的比较、like/contains/match、in/not in 等
      * 其他命令中的赋值（如 span="10m" 等）
    - 避免误报：允许函数对字段名的单引号（如 avg('field')）、by/fields/as 后的字段名单引号、作为左值的单引号字段名（'field'=...）。
    出现违规时抛出 ValueError 并提供修复建议片段。
    """
    try:
        violations = []

        # 1) 全局：任何 key='value' 的赋值均违规
        for m in re.finditer(r"\b[\w.]+\s*=\s*'[^']*'", spl):
            frag = m.group(0)
            corrected = re.sub(r"=\s*'([^']*)'", r'="\1"', frag)
            violations.append((frag, corrected))

        # 2) where 子句：函数与比较操作的值不得为单引号
        for segment in spl.split("|"):
            seg = segment.strip()
            if not seg:
                continue
            is_where = seg.lower().startswith("where")
            expr = seg[len("where"):].strip() if is_where else seg

            # like/contains/match 第二个参数不应为单引号
            for pattern in [
                r"\blike\s*\(\s*[^,]+,\s*'[^']*'\s*\)",
                r"\bcontains\s*\(\s*[^,]+,\s*'[^']*'\s*\)",
                r"\bmatch\s*\(\s*[^,]+,\s*'[^']*'\s*\)"
            ]:
                for m in re.finditer(pattern, expr, flags=re.IGNORECASE):
                    frag = m.group(0)
                    # 仅替换第二个参数的引号，保留字段名的单引号
                    corrected = re.sub(
                        r"\b(like|contains|match)\s*\(\s*([^,]+),\s*'([^']*)'\s*\)",
                        r'\1(\2, "\3")',
                        frag,
                        flags=re.IGNORECASE,
                    )
                    violations.append((frag, corrected))

            # in/not in 列表值不得为单引号
            for m in re.finditer(r"\b(?:in|not\s+in)\s*\(\s*'[^']*'(?:\s*,\s*'[^']*')*\s*\)", expr, flags=re.IGNORECASE):
                frag = m.group(0)
                corrected = re.sub(r"'([^']*)'", r'"\1"', frag)
                violations.append((frag, corrected))

            # 等值/不等比较右值不得为单引号
            for m in re.finditer(r"(?:=|!=|<>|>=|<=|=~|!~|~)\s*'[^']*'", expr):
                frag = m.group(0)
                corrected = frag.replace("'", '"')
                violations.append((frag, corrected))

        # 3) search/search2 主段：裸单引号文本（自由文本值）不得使用单引号
        first_segment = spl.split("|")[0].strip()
        if first_segment.lower().startswith(("search2", "search")):
            for m in re.finditer(r"'[^']*'", first_segment):
                frag = m.group(0)
                start, end = m.start(), m.end()
                # 允许的上下文：函数参数、by/fields/as 后的字段名、作为左值的字段名（后面紧跟=）
                prefix = first_segment[max(0, start-30):start]
                suffix = first_segment[end:min(len(first_segment), end+10)]

                allowed_prefix_keywords = ["by", "fields", "as"]
                if any(re.search(rf"\b{kw}\s*$", prefix) for kw in allowed_prefix_keywords):
                    continue
                if re.search(r"\([\s]*$", prefix):  # 函数左括号紧邻
                    continue
                if re.match(r"^\s*=", suffix):  # 左值场景：'field'=
                    continue
                # 允许作为操作符左操作数：'field' in / not in / 比较运算符
                if re.match(r"^\s*(in|not\s+in|=|!=|<>|>=|<=|=~|!~|~)\b", suffix, flags=re.IGNORECASE):
                    continue

                # 其他场景视为自由文本值，应改为双引号
                corrected = frag.replace("'", '"')
                violations.append((frag, corrected))

        # 4) 管道规则：where 等算子必须通过管道符引入（缺少时给出修复建议）
        try:
            stage_cmds = ["where"]
            for cmd in stage_cmds:
                if re.search(rf"\b{cmd}\b", first_segment, flags=re.IGNORECASE) and not re.match(rf"^\s*{cmd}\b", first_segment, flags=re.IGNORECASE):
                    for m in re.finditer(rf"\b{cmd}\b", first_segment, flags=re.IGNORECASE):
                        frag = first_segment[max(0, m.start()-30):min(len(first_segment), m.end()+30)]
                        corrected = first_segment[:m.start()].rstrip() + " | " + first_segment[m.start():]
                        violations.append((frag, corrected))
        except Exception:
            pass

        if violations:
            lines = [
                "SPL语法错误：字段值必须使用双引号或 where 等算子前缺少管道符",
                "建议修复片段："
            ]
            for frag, corrected in violations[:10]:
                lines.append(f"- {frag} -> {corrected}")
            if len(violations) > 10:
                lines.append(f"... 共发现 {len(violations)} 处")
            lines.append("规则：字段名用单引号，字段值用双引号（或三双引号）；where 等算子必须通过管道符 '|' 引入。")
            raise ValueError("\n".join(lines))
    except ValueError:
        # 违反规则时应向上抛出，让调用方中断执行并显示错误
        raise
    except Exception as e:
        # 校验实现异常时不影响正常查询，但记录到控制台便于排查
        console.print(f"[yellow]SPL语法校验异常：{str(e)}[/yellow]")



@function_registry.register(
    name="search_data",
    description=f"在KetaDB中搜索指标数据，支持SPL查询语言。请先使用`get_docs `函数获取SPL语法参考文档。",
    parameters={
        "type": "object",
        "properties": {
            "spl": {
                "type": "string",
                "description": "SPL查询语句，必须符合SPL语法规范。"
            },
            "limit": {
                "type": "integer",
                "description": "返回结果数量限制",
                "default": 100
            },
            "format_type": {
                "type": "string",
                "description": "输出格式 (text, json, csv, toon)",
                "default": "toon"
            },
            "output_file": {
                "type": "string",
                "description": "输出文件路径；提供则将结果写入文件"
            },
        },
        "required": ["spl"]
    }
)
def search_data(spl: str, limit: int = 100, format_type: str = "toon", output_file: Optional[str] = None) -> str:
    """搜索KetaDB中的数据"""

    try:
        spl = preflight_spl_fix(spl)
    except Exception:
        pass
    try:
        validate_mstats_syntax(spl)
    except ValueError as ve:
        return str(ve)
    except Exception:
        pass
    
    # 先执行查询；仅在结果为空或查询/解析出错时触发语法校验并返回错误信息
    try:
        resp = search_spl(spl=spl, limit=limit)
    except Exception as se:
        # 查询失败时，尝试语法校验；若存在引号违规则返回更友好的错误提示
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"搜索失败: {str(se)}, 请调用`get_docs log_search_syntax`函数获取SPL语法参考文档以确认查询语句是否正确。"

    # 构造结果输出对象可能抛错（返回结构异常等），同样走校验兜底
    try:
        result_output = search_result_output(resp)
    except Exception as pe:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"结果解析失败: {str(pe)}"

    # 正常格式化输出（有数据时直接返回结果）
    if format_type not in ["json", "csv", "text", "toon"]:
        format_type = "toon"

    text = result_output.get_formatted_string(format_type)

    # 结果为空时才做语法校验，提示可能的引号问题；无违规则返回空结果文本
    try:
        rows_count = len(getattr(result_output, "rows", []) or [])
    except Exception:
        rows_count = 0
    if rows_count == 0:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        # 返回友好提示：建议放宽/减少查询条件
        return (
            "未查询到数据。建议减少或放宽查询条件：尝试扩大时间范围（调整 `start`）、减少 搜索条件过滤项等！"
        )
    # 写文件输出
    if output_file and isinstance(output_file, str) and len(output_file.strip()) > 0:
        try:
            dir_name = os.path.dirname(output_file)
            if dir_name:
                os.makedirs(dir_name, exist_ok=True)
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(text)
            return f"[green]搜索结果已写入: {output_file}[/green]"
        except Exception as fe:
            return f"写入文件失败: {str(fe)}"

    if rows_count == limit:
        text += "\n注意：返回数量等于limit，可能还有更多数据，建议调整limit参数获取更多数据或者通过span等参数调整采样"
    return text


@function_registry.register(
    name="search_metadata",
    description="获取SPL查询的元数据信息",
    parameters={
        "type": "object",
        "properties": {
            "spl": {
                "type": "string",
                "description": "SPL查询语句"
            },
            "output_file": {
                "type": "string",
                "description": "输出文件路径；提供则将结果写入文件"
            }
        },
        "required": ["spl"]
    }
)
def search_metadata(spl: str, output_file: Optional[str] = None) -> str:
    """获取SPL查询的元数据"""
    try:
        if "show" in spl:
            return "show 命令不支持获取元数据，请直接使用search_data_for_metric函数获取指标元数据"
        meta = search_spl_meta(spl)
        table = get_asset_output(resp=meta)
        text = format_table(table, format="toon")
        if output_file and isinstance(output_file, str) and len(output_file.strip()) > 0:
            try:
                dir_name = os.path.dirname(output_file)
                if dir_name:
                    os.makedirs(dir_name, exist_ok=True)
                with open(output_file, "w", encoding="utf-8") as f:
                    f.write(text)
                console.print(f"[green]元数据(TOON)已写入: {output_file}[/green]")
            except Exception as fe:
                console.print(f"[yellow]写入文件失败: {str(fe)}[/yellow]")
        return text
    except Exception as e:
        return f"获取元数据失败: {str(e)}"


@function_registry.register(
    name="get_repo_fields",
    description=f"""获取指定仓库的字段信息和数据摘要。""",
    parameters={
        "type": "object",
        "properties": {
            "repository": {
                "type": "string",
                "description": "目标仓库名称"
            },
            "limit": {
                "type": "integer",
                "description": "返回的字段数量",
                "default": 10
            }
        },
        "required": ["repository"]
    }
)
def get_repo_fields(repository: str, limit: int = 20) -> str:
    """获取指定仓库的字段信息"""
    get_fields_spl = f"search2 repo=\"{repository}\""
    
    try:
        # 构建SPL查询，将时间参数直接嵌入到查询中
        # 调用search_summary获取字段信息，不传递时间参数（因为已经在SPL中了）
        fields = list(search_summary(spl=get_fields_spl).keys())
        # return fields
        # 格式化返回结果
        field_values = {}
        # 初始化计数器和有效字段计数器
        n = 0
        valid_fields_count = 0
        f = open("repo_fields.txt", "a", encoding="utf-8")
        # 遍历所有字段
        while n < len(fields) and valid_fields_count < limit:
            field = fields[n]
            n += 1
            f.write(f"{field}\n")
            
            # 跳过系统字段
            if field in ("_indexTime", "_time", "_raw", "_signature_id", "linenum", "repo"):
                continue
                
            get_fields_value_spl = f"""search2 start="-1h" repo="{repository}"
| where isNotNull('{field}')    
| stats count() as countKey1 by '{field}' with count() as total, distinct('{field}', 0) as distinct
| sort 3 by countKey1
| eval value='{field}', key="{field}"
| fields key, value, countKey1
"""
            rsp = search_spl(spl=get_fields_value_spl)["rows"]
            if not rsp:
                continue
            for item in rsp:
                if item[0] not in field_values:
                    field_values[item[0]] = {}
                field_values[item[0]][item[1]] = str(item[2])[:50] 
            # 增加有效字段计数
            valid_fields_count += 1
        
        rows = []
        rows.append(["repository", repository, None])
        rows.append(["total_fields", len(fields), None])
        rows.append(["return_fields", limit, None])
        for field_name, values in field_values.items():
            for val, cnt in values.items():
                rows.append([field_name, val, cnt])
        table = make_table(["field", "value", "count"], rows)
        text = table.get_formatted_string("toon")
        return text
        
    except Exception as e:
        import traceback
        return f"获取仓库字段信息失败: {str(e)}\n{traceback.format_exc()}"
def preflight_spl_fix(spl: str) -> str:
    try:
        s = spl or ""
        s = s.replace("｜", "|")
        s = re.sub(r"\b([\w.]+)\s*=\s*'([^']*)'", r'\1="\2"', s)
        first_segment = s.split("|")[0]
        if re.search(r"\bwhere\b", first_segment, flags=re.IGNORECASE) and not re.match(r"^\s*where\b", first_segment, flags=re.IGNORECASE):
            s = re.sub(r"(.*?)(\bwhere\b)", lambda m: m.group(1).rstrip()+" | "+m.group(2), first_segment, flags=re.IGNORECASE) + ("|"+"|".join(s.split("|")[1:]) if "|" in s else "")
        s = re.sub(r"\b(?:in|not\s+in)\s*\(([^)]*)\)", lambda m: "in("+re.sub(r"'([^']*)'", r'"\\1"', m.group(1))+")" if m.group(0).lower().startswith("in") else "not in("+re.sub(r"'([^']*)'", r'"\\1"', m.group(1))+")", s, flags=re.IGNORECASE)
        return s
    except Exception:
        return spl

def validate_mstats_syntax(spl: str) -> None:
    t = (spl or "").strip()
    tl = t.lower()
    if not tl.startswith("mstats"):
        if "mstats" in tl:
            pass
        else:
            return
    if re.search(r"\bsearch2\b", tl):
        raise ValueError("SPL语法错误：指标查询不得与search2混用，请使用mstats并在开头给出时间参数")
    if not re.match(r"^\s*mstats\s+", tl):
        raise ValueError("SPL语法错误：mstats必须作为查询开头")
    if not re.search(r"mstats\s+start=|mstats\s+end=", tl):
        raise ValueError("SPL语法错误：mstats后需紧跟时间参数start/end")
    for m in re.finditer(r"\b(avg|sum|min|max|rate|topseries|hperc|distinct)\b", tl):
        pass
    if re.search(r"\bsort\s+by\b", tl) is None and re.search(r"\btopseries\b", tl):
        raise ValueError("SPL语法错误：使用topseries时建议配合sort by 指定排序字段")


@function_registry.register(
    name="simple_search",
    description="根据参数组装日志SPL并搜索数据。支持仓库、过滤条件、排序与限制。",
    parameters={
        "type": "object",
        "properties": {
            "repository": {"type": "string", "description": "仓库名"},
            "start": {"type": "string", "description": "开始时间，如 -1h、-d@d"},
            "end": {"type": "string", "description": "结束时间，如 @h、@d"},
            "filters": {
                "type": "array",
                "description": "过滤条件列表，支持与或非",
                "items": {
                    "type": "object",
                    "properties": {
                        "field": {"type": "string", "description": "字段名"},
                        "op": {"type": "string", "description": "操作符，如 =, !=, >=, <=, >, <, in, not in"},
                        "value": {"description": "值(非in时)", "anyOf": [{"type": "string"}, {"type": "number"}]},
                        "values": {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "number"}]}},
                        "text": {"type": "string", "description": "全文匹配文本"},
                        "not": {"type": "boolean", "description": "取反"}
                    }
                }
            },
            "stats": {
                "type": "array",
                "description": "聚合算子列表，如 count()/sum('f') 等",
                "items": {
                    "type": "object",
                    "properties": {
                        "func": {"type": "string", "description": "聚合函数，如 count,sum,avg,min,max,distinct,percentile"},
                        "field": {"type": "string", "description": "聚合字段（count可省略）"},
                        "alias": {"type": "string", "description": "别名"},
                        "args": {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "number"}]}, "description": "额外参数，如percentile的百分位"}
                    }
                }
            },
            "group_by": {
                "type": "array",
                "items": {"type": "string"},
                "description": "分组字段列表"
            },
            "logic": {"type": "string", "enum": ["AND", "OR"], "default": "AND", "description": "过滤条件组合逻辑"},
            "sort_field": {"type": "string", "description": "排序字段"},
            "sort_order": {"type": "string", "enum": ["asc", "desc"], "default": "desc", "description": "排序顺序"},
            "limit": {"type": "integer", "default": 100, "description": "数量限制"},
            "format_type": {"type": "string", "default": "toon", "description": "输出格式(text,json,csv,toon)，优先选择toon"},
            "output_file": {"type": "string", "description": "写入文件路径"}
        },
        "required": ["repository"]
    }
)
def simple_search(
    repository: str,
    start: Optional[str] = None,
    end: Optional[str] = None,
    filters: Optional[list] = None,
    stats: Optional[list] = None,
    group_by: Optional[list] = None,
    logic: str = "AND",
    sort_field: Optional[str] = None,
    sort_order: str = "desc",
    limit: int = 100,
    format_type: str = "toon",
    output_file: Optional[str] = None,
) -> str:
    def is_number(v: Union[str, int, float]) -> bool:
        if isinstance(v, (int, float)):
            return True
        if isinstance(v, str):
            return re.fullmatch(r"-?\d+(?:\.\d+)?", v) is not None
        return False

    def quote_val(v: Union[str, int, float]) -> str:
        if is_number(v):
            return str(v)
        s = str(v)
        if '"' in s:
            return f'"""{s}"""'
        return f'"{s}"'

    def build_condition(fobj: dict) -> str:
        if not isinstance(fobj, dict):
            return ""
        negate = bool(fobj.get("not"))
        if fobj.get("text") and not fobj.get("field"):
            expr = quote_val(fobj.get("text"))
            return f"NOT ({expr})" if negate else expr
        field = fobj.get("field")
        op = (fobj.get("op") or "=").strip().lower()
        if not field:
            return ""
        l = f"'{field}'"
        if op in ("in", "not in"):
            vals = fobj.get("values") or []
            parts = []
            for vv in vals:
                parts.append(quote_val(vv))
            expr = f"{l} {op} (" + ",".join(parts) + ")"
            return f"NOT ({expr})" if negate else expr
        if op in ("contains", "like"):
            tv = fobj.get("value")
            if tv is None:
                return ""
            patt = str(tv)
            if op == "contains" and not (patt.startswith("*") or patt.endswith("*")):
                patt = f"*{patt}*"
            expr = quote_val(patt)
            return f"NOT ({expr})" if negate else expr
        rv = fobj.get("value")
        if rv is None:
            return ""
        expr = f"{l}{op if op in ['=','!=','>=','<=','>','<','=~','!~','~'] else '='}{quote_val(rv)}"
        return f"NOT ({expr})" if negate else expr

    segments = ["search2"]
    time_parts = []
    if isinstance(start, str) and start.strip():
        time_parts.append(f'start="{start.strip()}"')
    if isinstance(end, str) and end.strip():
        time_parts.append(f'end="{end.strip()}"')
    if time_parts:
        segments.append(" ".join(time_parts))
    segments.append(f'repo="{repository}"')

    conds = []
    if isinstance(filters, list) and filters:
        for it in filters:
            c = build_condition(it)
            if c:
                conds.append(c)
    if conds:
        joiner = f" {('OR' if (logic or 'AND').upper() == 'OR' else 'AND')} "
        segments.append("(" + joiner.join(conds) + ")")

    spl = " ".join(segments)
    pipes = []
    agg_parts = []
    if isinstance(stats, list) and stats:
        for agg in stats:
            if not isinstance(agg, dict):
                continue
            func = str(agg.get("func") or "").strip()
            if not func:
                continue
            fl = str(agg.get("field") or "").strip()
            args = agg.get("args") or []
            func_l = func.lower()
            if func_l == "count":
                expr = "count()"
            else:
                if not fl and func_l not in ("earliest", "latest"):
                    continue
                if args and isinstance(args, list) and len(args) > 0:
                    arg_strs = []
                    for a in args:
                        if isinstance(a, (int, float)) or (isinstance(a, str) and re.fullmatch(r"-?\d+(?:\.\d+)?", a)):
                            arg_strs.append(str(a))
                        else:
                            s = str(a)
                            arg_strs.append(f'"{s}"' if '"' not in s else f'"""{s}"""')
                    expr = f"{func}('{fl}', " + ",".join(arg_strs) + ")"
                else:
                    if fl:
                        expr = f"{func}('{fl}')"
                    else:
                        expr = f"{func}()"
            alias = str(agg.get("alias") or "").strip()
            if alias:
                expr = expr + f" as '{alias}'"
            agg_parts.append(expr)
    if agg_parts:
        by_part = ""
        if isinstance(group_by, list) and group_by:
            gb = [f"'{str(g).strip()}'" for g in group_by if str(g).strip()]
            if gb:
                by_part = " by " + ", ".join(gb)
        pipes.append("stats " + ", ".join(agg_parts) + by_part)
    if sort_field and isinstance(sort_field, str) and sort_field.strip():
        order = (sort_order or "desc").lower()
        if order not in ("asc", "desc"):
            order = "desc"
        pipes.append(f"sort by '{sort_field.strip()}' {order}")
    pipes.append(f"limit {int(limit) if isinstance(limit, int) else 100}")
    if pipes:
        spl = spl + " | " + " | ".join(pipes)

    try:
        spl = preflight_spl_fix(spl)
    except Exception:
        pass
    try:
        validate_spl_value_quotes(spl)
    except ValueError as ve:
        return str(ve)
    except Exception:
        pass

    logger.info(f"生成的执行SPL: {spl}")

    try:
        resp = search_spl(spl=spl, limit=limit)
    except Exception as se:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"搜索失败: {str(se)}, 请调用`get_docs log_search_syntax`函数获取SPL语法参考文档以确认查询语句是否正确。"

    try:
        result_output = search_result_output(resp)
    except Exception as pe:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"结果解析失败: {str(pe)}"

    if format_type not in ["json", "csv", "text", "toon"]:
        format_type = "toon"
    text = result_output.get_formatted_string(format_type)
    try:
        rows_count = len(getattr(result_output, "rows", []) or [])
    except Exception:
        rows_count = 0
    if rows_count == 0:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return (
            "未查询到数据。建议减少或放宽查询条件：尝试扩大时间范围（调整 `start`）、减少 搜索条件过滤项等！"
        )
    if output_file and isinstance(output_file, str) and len(output_file.strip()) > 0:
        try:
            dir_name = os.path.dirname(output_file)
            if dir_name:
                os.makedirs(dir_name, exist_ok=True)
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(text)
            return f"搜索结果已写入: {output_file}"
        except Exception as fe:
            return f"写入文件失败: {str(fe)}"
    if rows_count == limit:
        text += "\n注意：返回数量等于limit，可能还有更多数据，建议调整limit参数获取更多数据或者通过span等参数调整采样"
    return text


@function_registry.register(
    name="search_traces_by_traceid",
    description="根据traceid查询TRACES链路数据，支持指定返回字段。",
    parameters={
        "type": "object",
        "properties": {
            "traceid": {"type": "string", "description": "链路ID"},
            "start": {"type": "string", "description": "开始时间", "default": "-3h"},
            "return_fields": {
                "type": "array",
                "description": "需要返回的字段列表",
                "items": {"type": "string"}
            },
            "limit": {"type": "integer", "description": "返回数量限制", "default": 100},
            "format_type": {"type": "string", "description": "输出格式(text,json,csv,toon)", "default": "toon"},
            "output_file": {"type": "string", "description": "输出文件路径"}
        },
        "required": ["traceid"]
    }
)
def search_traces_by_traceid(
    traceid: str,
    start: str = "-3h",
    return_fields: Optional[list] = None,
    limit: int = 100,
    format_type: str = "toon",
    output_file: Optional[str] = None,
) -> str:
    segments = [
        "search2",
        f'start="{start.strip()}"' if isinstance(start, str) and start.strip() else 'start="-3h"',
        '_repo_type="TRACES"',
        f'(\'traceId\'="*{traceid}*")'
    ]
    spl = " ".join(segments)
    pipes = [
        "eval duration = durationNanos / 1000000",
        'eval attrs = arr_join(attrs, "[:]"), resource.attrs = arr_join(resource.attrs, "[:]")',
    ]
    if isinstance(return_fields, list) and len(return_fields) > 0:
        fs = ", ".join([str(f).strip() for f in return_fields if str(f).strip()])
        if fs:
            pipes.append(f"fields {fs}")
    pipes.append(f"limit {int(limit) if isinstance(limit, int) else 100}")
    spl = spl + " | " + " | ".join(pipes)

    try:
        spl = preflight_spl_fix(spl)
    except Exception:
        pass
    try:
        validate_spl_value_quotes(spl)
    except ValueError as ve:
        return str(ve)
    except Exception:
        pass

    try:
        resp = search_spl(spl=spl, limit=limit)
    except Exception as se:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"搜索失败: {str(se)}"

    try:
        result_output = search_result_output(resp)
    except Exception as pe:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"结果解析失败: {str(pe)}"

    if format_type not in ["json", "csv", "text", "toon"]:
        format_type = "toon"
    text = result_output.get_formatted_string(format_type)
    try:
        rows_count = len(getattr(result_output, "rows", []) or [])
    except Exception:
        rows_count = 0
    if rows_count == 0:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return "未查询到数据。请确认traceid正确，并尝试扩大时间范围。"
    if output_file and isinstance(output_file, str) and len(output_file.strip()) > 0:
        try:
            dir_name = os.path.dirname(output_file)
            if dir_name:
                os.makedirs(dir_name, exist_ok=True)
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(text)
            return f"搜索结果已写入: {output_file}"
        except Exception as fe:
            return f"写入文件失败: {str(fe)}"
    if rows_count == limit:
        text += "\n注意：返回数量等于limit，可能还有更多数据，可调整limit或时间范围。"
    return text


@function_registry.register(
    name="search_trace_exceptions_by_traceid",
    description="根据traceid查询TRACES中的异常事件详情，支持指定返回字段。",
    parameters={
        "type": "object",
        "properties": {
            "traceid": {"type": "string", "description": "链路ID"},
            "start": {"type": "string", "description": "开始时间", "default": "-3h"},
            "return_fields": {
                "type": "array",
                "description": "需要返回的字段列表",
                "items": {"type": "string"}
            },
            "limit": {"type": "integer", "description": "返回数量限制", "default": 100},
            "format_type": {"type": "string", "description": "输出格式(text,json,csv,toon)", "default": "toon"},
            "output_file": {"type": "string", "description": "输出文件路径"}
        },
        "required": ["traceid"]
    }
)
def search_trace_exceptions_by_traceid(
    traceid: str,
    start: str = "-3h",
    return_fields: Optional[list] = None,
    limit: int = 100,
    format_type: str = "toon",
    output_file: Optional[str] = None,
) -> str:
    segments = [
        "search2",
        f'start="{start.strip()}"' if isinstance(start, str) and start.strip() else 'start="-3h"',
        '_repo_type="TRACES"',
        f'(\'traceId\'="*{traceid}*")'
    ]
    spl = " ".join(segments)
    pipes = [
        "mvexpand events",
        'eval events.name = jsonpath(events, "name")',
        'where events.name = "exception"',
        'eval exception.type = jsonpath(events, """attributes[\"exception.type\"]"""), exception.message = jsonpath(events, """attributes[\"exception.message\"]"""), exception.stacktrace = jsonpath(events, """attributes[\"exception.stacktrace\"]"""), exception.escaped = jsonpath(events, """attributes[\"exception.escaped\"]""")',
        'eval timeUnixNano = jsonpath(events, "timeUnixNano"), _time = timeUnixNano / 1000000',
        'eval _raw = if (exception.message != "", concat(exception.type, ": ", exception.message), exception.type)',
        'rename kind as span.kind, name as span.name',
        "fields _time, _raw, timeUnixNano, env, service, servicetype, version, instance, host, scope.name, scope.version, span.kind, span.name, traceId, spanId, parentSpanId, exception.type, exception.message, exception.stacktrace, exception.escaped"
    ]
    if isinstance(return_fields, list) and len(return_fields) > 0:
        fs = ", ".join([str(f).strip() for f in return_fields if str(f).strip()])
        if fs:
            pipes.append(f"fields {fs}")
    pipes.append(f"limit {int(limit) if isinstance(limit, int) else 100}")
    spl = spl + " | " + " | ".join(pipes)

    try:
        spl = preflight_spl_fix(spl)
    except Exception:
        pass
    try:
        validate_spl_value_quotes(spl)
    except ValueError as ve:
        return str(ve)
    except Exception:
        pass

    try:
        resp = search_spl(spl=spl, limit=limit)
    except Exception as se:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"搜索失败: {str(se)}"

    try:
        result_output = search_result_output(resp)
    except Exception as pe:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return f"结果解析失败: {str(pe)}"

    if format_type not in ["json", "csv", "text", "toon"]:
        format_type = "toon"
    text = result_output.get_formatted_string(format_type)
    try:
        rows_count = len(getattr(result_output, "rows", []) or [])
    except Exception:
        rows_count = 0
    if rows_count == 0:
        try:
            validate_spl_value_quotes(spl)
        except ValueError as ve:
            return str(ve)
        return "未查询到异常事件。请确认traceid正确，并尝试扩大时间范围。"
    if output_file and isinstance(output_file, str) and len(output_file.strip()) > 0:
        try:
            dir_name = os.path.dirname(output_file)
            if dir_name:
                os.makedirs(dir_name, exist_ok=True)
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(text)
            return f"搜索结果已写入: {output_file}"
        except Exception as fe:
            return f"写入文件失败: {str(fe)}"
    if rows_count == limit:
        text += "\n注意：返回数量等于limit，可能还有更多异常事件，可调整limit或时间范围。"
    return text


@function_registry.register(
    name="simple_metric_search",
    description="根据指标名称与聚合方式生成 mstats 查询并返回结果。",
    parameters={
        "type": "object",
        "properties": {
            "metrics": {
                "type": "array",
                "description": "指标聚合列表",
                "items": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string", "description": "指标名称"},
                        "func": {"type": "string", "description": "聚合函数，如 avg,sum,min,max,distinct,percentile,hperc,sumRate"},
                        "alias": {"type": "string", "description": "别名"},
                        "args": {"type": "array", "items": {"anyOf": [{"type": "number"}, {"type": "string"}]}, "description": "函数额外参数"}
                    },
                    "required": ["name", "func"]
                }
            },
            "start": {"type": "string", "description": "开始时间，如 -1h、-d@d", "default": "-1h"},
            "end": {"type": "string", "description": "结束时间，如 @h、@d"},
            "span": {"type": "string", "description": "时间粒度，如 1m、5m"},
            "timeshift": {"type": "string", "description": "时间偏移，如 -1h"},
            "filters": {
                "type": "array",
                "description": "标签过滤条件（用于 where 段）",
                "items": {
                    "type": "object",
                    "properties": {
                        "field": {"type": "string"},
                        "op": {"type": "string", "description": "=, !=, like, in, not in"},
                        "value": {"anyOf": [{"type": "string"}, {"type": "number"}]},
                        "values": {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "number"}]}},
                        "not": {"type": "boolean"}
                    }
                }
            },
            "logic": {"type": "string", "enum": ["AND", "OR"], "default": "AND"},
            "group_by": {"type": "array", "items": {"type": "string"}},
            "sort_field": {"type": "string"},
            "sort_order": {"type": "string", "enum": ["asc", "desc"], "default": "desc"},
            "limit": {"type": "integer", "default": 100},
            "format_type": {"type": "string", "default": "toon", "description": "输出格式(text,json,csv,toon)"},
            "output_file": {"type": "string"}
        },
        "required": ["metrics"]
    }
)
def simple_metric_search(
    metrics: list,
    start: Optional[str] = "-1h",
    end: Optional[str] = None,
    span: Optional[str] = None,
    timeshift: Optional[str] = None,
    filters: Optional[list] = None,
    logic: str = "AND",
    group_by: Optional[list] = None,
    sort_field: Optional[str] = None,
    sort_order: str = "desc",
    limit: int = 100,
    format_type: str = "toon",
    output_file: Optional[str] = None,
) -> str:
    spl = build_mstats_spl(
        metrics=metrics,
        start=start,
        end=end,
        span=span,
        timeshift=timeshift,
        filters=filters,
        logic=logic,
        group_by=group_by,
        sort_field=sort_field,
        sort_order=sort_order,
        limit=limit,
    )

    try:
        spl = preflight_spl_fix(spl)
    except Exception:
        pass
    try:
        validate_mstats_syntax(spl)
    except ValueError as ve:
        return str(ve)
    except Exception:
        pass

    logger.info(f"生成的指标SPL: {spl}")

    try:
        resp = search_spl(spl=spl, limit=limit)
    except Exception as se:
        return f"搜索失败: {str(se)}"
    try:
        result_output = search_result_output(resp)
    except Exception as pe:
        return f"结果解析失败: {str(pe)}"
    if format_type not in ["json", "csv", "text", "toon"]:
        format_type = "toon"
    text = result_output.get_formatted_string(format_type)
    try:
        rows_count = len(getattr(result_output, "rows", []) or [])
    except Exception:
        rows_count = 0
    if rows_count == 0:
        return (
            "未查询到数据。建议调整时间范围（start/end/span）、检查分组字段是否为标签，或放宽过滤条件。SPL: " + spl
        )
    if output_file and isinstance(output_file, str) and len(output_file.strip()) > 0:
        try:
            dir_name = os.path.dirname(output_file)
            if dir_name:
                os.makedirs(dir_name, exist_ok=True)
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(text)
            return f"[green]搜索结果已写入: {output_file}[/green]"
        except Exception as fe:
            return f"[yellow]写入文件失败: {str(fe)}[/yellow]"
    if rows_count == limit:
        text += "\n注意：返回数量等于limit，可能还有更多数据，建议调整limit或span参数。"
    return text

def build_mstats_spl(
    metrics: List[Dict],
    start: Optional[str] = None,
    end: Optional[str] = None,
    span: Optional[str] = None,
    timeshift: Optional[str] = None,
    filters: Optional[List[Dict]] = None,
    logic: str = "AND",
    group_by: Optional[List[str]] = None,
    sort_field: Optional[str] = None,
    sort_order: str = "desc",
    limit: int = 100,
) -> str:
    def is_number(v: Union[str, int, float]) -> bool:
        if isinstance(v, (int, float)):
            return True
        if isinstance(v, str):
            return re.fullmatch(r"-?\d+(?:\.\d+)?", v) is not None
        return False
    def quote_val(v: Union[str, int, float]) -> str:
        if is_number(v):
            return str(v)
        s = str(v)
        return f'"""{s}"""' if '"' in s else f'"{s}"'
    def build_where_condition(c: dict) -> str:
        if not isinstance(c, dict):
            return ""
        field = str(c.get("field") or "").strip()
        if not field:
            return ""
        op = (c.get("op") or "=").strip().lower()
        negate = bool(c.get("not"))
        l = f"'{field}'"
        if op in ("in", "not in"):
            vals = c.get("values") or []
            parts = [f"{l}={quote_val(v)}" for v in vals]
            expr = "(" + " OR ".join(parts) + ")"
            return f"NOT {expr}" if negate or op == "not in" else expr
        if op == "like":
            v = c.get("value")
            if v is None:
                return ""
            expr = f"like({l}, {quote_val(v)})"
            return f"NOT ({expr})" if negate else expr
        v = c.get("value")
        if v is None:
            return ""
        expr = f"{l}{op if op in ['=','!='] else '='}{quote_val(v)}"
        return f"NOT ({expr})" if negate else expr
    segs = ["mstats"]
    time_params = []
    if isinstance(start, str) and start.strip():
        time_params.append(f'start="{start.strip()}"')
    if isinstance(end, str) and end.strip():
        time_params.append(f'end="{end.strip()}"')
    if isinstance(span, str) and span.strip():
        time_params.append(f'span="{span.strip()}"')
    if isinstance(timeshift, str) and timeshift.strip():
        time_params.append(f'timeshift="{timeshift.strip()}"')
    if time_params:
        segs.append(" ".join(time_params))
    agg_exprs = []
    for m in metrics or []:
        if not isinstance(m, dict):
            continue
        name = str(m.get("name") or "").strip()
        func = str(m.get("func") or "").strip()
        if not name or not func:
            continue
        args = m.get("args") or []
        arg_str = ""
        if args:
            arg_parts = []
            for a in args:
                if is_number(a):
                    arg_parts.append(str(a))
                else:
                    s = str(a)
                    arg_parts.append(f'"{s}"' if '"' not in s else f'"""{s}"""')
            arg_str = ", " + ",".join(arg_parts)
        expr = f"{func}('{name}'{arg_str})"
        alias = str(m.get("alias") or "").strip()
        if alias:
            expr = expr + f" as '{alias}'"
        agg_exprs.append(expr)
    by_part = ""
    if isinstance(group_by, list) and group_by:
        gb = [f"'{str(g).strip()}'" for g in group_by if str(g).strip()]
        if gb:
            by_part = " by " + ", ".join(gb)
    if agg_exprs:
        segs.append(" ".join([", ".join(agg_exprs) + by_part]))
    spl = " ".join(segs)
    pipes = []
    where_parts = []
    if isinstance(filters, list) and filters:
        for c in filters:
            wc = build_where_condition(c)
            if wc:
                where_parts.append(wc)
    if where_parts:
        joiner = f" {('OR' if (logic or 'AND').upper() == 'OR' else 'AND')} "
        pipes.append("where " + "(" + joiner.join(where_parts) + ")")
    if sort_field and isinstance(sort_field, str) and sort_field.strip():
        order = (sort_order or "desc").lower()
        if order not in ("asc", "desc"):
            order = "desc"
        pipes.append(f"sort by '{sort_field.strip()}' {order}")
    pipes.append(f"limit {int(limit) if isinstance(limit, int) else 100}")
    if pipes:
        spl = spl + " | " + " | ".join(pipes)
    return spl
