# -*- coding: utf-8 -*-
"""编排层：串联 数据库查询 -> LLM调用 -> 格式校验解析 -> 组装最终JSON。

单条记录的处理逻辑（classify_param_text）:
    1. 若新参数为空/None -> 直接判定无附件，不调用LLM(省成本，也没有可判断内容)。
    2. 调用 LLM 拿到调用结果(LLMCallResult)。
       - 若调用本身失败(网络/接口错误,重试耗尽) -> 记日志,标记 llm_call_error。
    3. 用 parser.parse_llm_output(mode=..., enable_other=...) 做严格格式校验。
       - 校验通过 -> 组装结果,记日志,结束。
       - 校验失败(FormatError) -> 记日志,重新调用LLM重试。
    4. 格式校验重试全部耗尽 -> 标记 status="llm_format_error"，不让脏数据进入最终结果。

无论上述哪一步，只要 OrderLogger 存在，每一次尝试都会被完整记录——包括发给
模型的完整对话、模型的原始回复（如果有）、格式校验是否通过。这是本模块与
日志模块的核心约定：不因为"这次调用失败了"就跳过记录，失败的调用恰恰最
需要被记下来供排查。

批量处理(classify_batch)在此基础上先按 (标识符, 键类型) 查数据库，回取出
{标识符: {param, 总排号}} 映射；对查不到的标识符直接标记 status="not_found"，
不发起LLM调用（也不生成日志，因为根本没有"新参数"可供判断，没有对话可记）。

enable_other 控制是否允许模型使用"其他"兜底类目（默认 False），从
classify_batch/classify_single 一路透传到 llm_client.classify_raw（决定
提示词要不要包含"其他"）和 parser.parse_llm_output（决定校验要不要放行
"其他"）。本模块是唯一同时持有这两个调用点的地方，因此由本模块负责把
同一个值传给两边，不暴露给更上层去分别设置从而导致不一致。
"""
from __future__ import annotations

import logging
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

from db import DatabaseError, fetch_params_by_ids
from llm_client import LLMCallResult, LLMClient
from order_logger import OrderLogger
from parser import FormatError, ParsedResult, parse_llm_output

logger = logging.getLogger(__name__)


def _empty_result(zong_pai_hao: str, status: str) -> dict[str, Any]:
    return {
        "zong_pai_hao": zong_pai_hao,
        "status": status,
        "has_attachment": None,
        "types": [],
        # 未发起 LLM 调用的结果（not_found/db_error），meta 给出一致的空壳，
        # 便于下游统一读取（elapsed_ms 为 None，统计时按 0 处理）。
        "meta": {"elapsed_ms": None, "attempts": 0, "llm": None},
    }


def _build_usage(u: dict[str, Any]) -> dict[str, Any]:
    """把原始 usage dict 整理成对外结构，并算出缓存率。

    缓存率 = prompt_cache_hit_tokens / prompt_tokens（prompt_tokens 为 0 时按
    0.0 处理，避免除零）。保留 4 位小数。
    """
    prompt = u.get("prompt_tokens", 0) or 0
    completion = u.get("completion_tokens", 0) or 0
    total = u.get("total_tokens", 0) or 0
    hit = u.get("prompt_cache_hit_tokens", 0) or 0
    miss = u.get("prompt_cache_miss_tokens", 0) or 0
    cache_hit_rate = round(hit / prompt, 4) if prompt else 0.0
    return {
        "prompt_tokens": prompt,
        "completion_tokens": completion,
        "total_tokens": total,
        "prompt_cache_hit_tokens": hit,
        "prompt_cache_miss_tokens": miss,
        "cache_hit_rate": cache_hit_rate,
    }


def _make_meta(
    call_result: LLMCallResult | None, start: float, attempt: int
) -> dict[str, Any]:
    """组装单条结果的 meta 对象。

    elapsed_ms 为从 classify_param_text 入口到此处的总耗时（含重试+解析）；
    若本次有成功调用，则 llm 块记录模型名、单次调用耗时与整理后的 usage；
    否则（调用失败/未调用）llm 为 None。
    """
    elapsed_ms = (time.perf_counter() - start) * 1000
    if call_result is not None and call_result.error is None and call_result.usage is not None:
        llm_block = {
            "model": call_result.model,
            "elapsed_ms": call_result.elapsed_ms,
            "usage": _build_usage(call_result.usage),
        }
    else:
        llm_block = None
    return {"elapsed_ms": elapsed_ms, "attempts": attempt, "llm": llm_block}


def classify_param_text(
    llm_client: LLMClient,
    param_text: str | None,
    mode: str = "coarse",
    format_retry: int = 2,
    order_logger: OrderLogger | None = None,
    enable_other: bool = False,
) -> tuple[str, bool | None, list[str], dict[str, Any]]:
    """对单条"新参数"文本做分类，返回 (status, has_attachment, types)。

    mode 决定分类粒度："coarse"只输出大类(资料/配件/耗材，启用 enable_other
    时还有"其他")；"fine"输出"大类:细分类目"（如"配件:针型阀"）的列表。
    不再有单独的 fine_types 返回值——fine 模式下细分类目已经内嵌在 types
    的每个元素里。

    enable_other 控制是否允许模型使用"其他"兜底那些不在既有枚举范围内的
    附件，默认关闭；关闭时的行为与不支持"其他"之前完全一致。

    status 取值:
        "ok"                - 正常识别成功
        "empty_param"       - 新参数字段为空，直接判无附件，未调用LLM
        "llm_call_error"    - LLM 网络/接口调用失败（重试耗尽）
        "llm_format_error"  - LLM 返回内容格式校验多次失败（重试耗尽）

    format_retry 控制"格式校验失败后重新请求LLM"的次数，与 LLMClient 内部的
    网络层重试是两回事：网络重试解决的是"请求没成功"，这里解决的是
    "请求成功了但模型没按格式输出"。两层重试互不影响，各自独立计数。

    order_logger 若提供，会记录本次分类过程中的每一次LLM往返（无论成败）。
    """
    start = time.perf_counter()

    if param_text is None or not str(param_text).strip():
        if order_logger is not None:
            order_logger.log_param_text(param_text)
            order_logger.log_final_result("empty_param", None)
        return "empty_param", False, [], _make_meta(None, start, 0)

    if order_logger is not None:
        order_logger.log_param_text(param_text)

    last_call_result: LLMCallResult | None = None
    last_format_err: Exception | None = None
    for attempt in range(1, format_retry + 1):
        call_result = llm_client.classify_raw(param_text, mode=mode, enable_other=enable_other)
        last_call_result = call_result

        if call_result.error is not None:
            logger.error("LLM 调用失败: %s", call_result.error)
            if order_logger is not None:
                order_logger.log_llm_attempt(attempt, call_result, None, None)
                order_logger.log_final_result("llm_call_error", None)
            return "llm_call_error", None, [], _make_meta(call_result, start, attempt)

        try:
            parsed = parse_llm_output(call_result.raw_response, mode=mode, enable_other=enable_other)
            if order_logger is not None:
                order_logger.log_llm_attempt(attempt, call_result, parsed, None)
                order_logger.log_final_result("ok", parsed)
            return "ok", parsed.has_attachment, parsed.types, _make_meta(call_result, start, attempt)
        except FormatError as e:
            last_format_err = e
            logger.warning(
                "LLM 输出格式校验失败 (第 %d/%d 次): %s | 原始输出: %r",
                attempt, format_retry, e, call_result.raw_response,
            )
            if order_logger is not None:
                order_logger.log_llm_attempt(attempt, call_result, None, str(e))

    logger.error("格式校验重试 %d 次后仍失败，放弃: %s", format_retry, last_format_err)
    if order_logger is not None:
        order_logger.log_final_result("llm_format_error", None)
    return "llm_format_error", None, [], _make_meta(last_call_result, start, format_retry)


def classify_batch(
    db_cfg: dict[str, Any],
    llm_cfg: dict[str, Any],
    id_list: list[tuple[str, str]],
    max_workers: int = 8,
    mode: str = "coarse",
    log_dir: str | Path | None = None,
    enable_other: bool = False,
) -> list[dict[str, Any]]:
    """批量分类入口：查库 -> 并发调用LLM -> 组装结果列表。

    id_list: list[(标识符, 键类型)]，键类型取值：
        "sn" -> 标识符为总排号，按总排号列查询；
        "id" -> 标识符为数据库真实 ID，按 id_field 列查询并回取总排号。
    两种键类型可在同一次调用中混用。

    log_dir 若提供，会为每个实际发起LLM调用的总排号在该目录下生成一个日志
    文件；查库失败(not_found)或参数为空(empty_param)的总排号不生成日志文件，
    因为它们本就没有与LLM的对话内容可记。

    返回结果顺序与输入 id_list 一致，即使某些查询失败或格式错误也会
    补全为对应 status 的占位结果，保证"输入N条，输出N条"。结果中的
    zong_pai_hao 一律为回查到的总排号（键类型为 id 时由数据库回取；键类型为
    sn 时即输入本身；not_found 时回退为输入标识符以便追溯）。
    """
    # 去重但保留顺序，避免用户传入重复标识符导致重复查询/调用
    seen: set[str] = set()
    unique_ids: list[tuple[str, str]] = []
    for item in id_list:
        idt = item[0]
        if idt not in seen:
            seen.add(idt)
            unique_ids.append(item)

    try:
        param_map = fetch_params_by_ids(db_cfg, unique_ids)
    except DatabaseError as e:
        logger.error("数据库查询失败，本批次全部标记为 db_error: %s", e)
        return [_empty_result(idt, "db_error") for idt, _ in id_list]

    llm_client = LLMClient(llm_cfg)
    results_by_id: dict[str, dict[str, Any]] = {}

    # 待分类项：(标识符, 对应总排号)；not_found(未查到) 直接占位
    ids_to_classify: list[tuple[str, str]] = []
    for idt, _ in unique_ids:
        entry = param_map.get(idt)
        if entry is not None and entry["param"] is not None:
            ids_to_classify.append((idt, entry["sn"]))
        else:
            # 未查到：zong_pai_hao 回退为输入标识符（可能是 ID 或 SN）
            results_by_id[idt] = _empty_result(idt, "not_found")

    def _worker(idt: str, sn: str | None) -> tuple[str, dict[str, Any]]:
        order_logger = OrderLogger(log_dir, sn, mode) if log_dir is not None else None
        status, has_attachment, types, meta = classify_param_text(
            llm_client, param_map[idt]["param"], mode=mode, order_logger=order_logger,
            enable_other=enable_other,
        )
        return idt, {
            "zong_pai_hao": sn,
            "status": status,
            "has_attachment": has_attachment,
            "types": types,
            "meta": meta,
        }

    if ids_to_classify:
        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            futures = [pool.submit(_worker, idt, sn) for idt, sn in ids_to_classify]
            for fut in as_completed(futures):
                idt, result = fut.result()
                results_by_id[idt] = result

    # 按原始输入顺序（含重复项）展开最终结果
    return [results_by_id[idt] for idt, _ in id_list]


def classify_single(
    db_cfg: dict[str, Any],
    llm_cfg: dict[str, Any],
    identifier: str,
    key: str = "sn",
    mode: str = "coarse",
    log_dir: str | Path | None = None,
    enable_other: bool = False,
) -> dict[str, Any]:
    """单个总排号/ID 分类的便捷封装。key 取值 "sn"(总排号) 或 "id"(真实ID)。"""
    return classify_batch(
        db_cfg, llm_cfg, [(identifier, key)], max_workers=1, mode=mode, log_dir=log_dir,
        enable_other=enable_other,
    )[0]


def summarize_results(
    results: list[dict[str, Any]],
    wall_clock_ms: float | None = None,
) -> dict[str, Any]:
    """汇总一批 classify_batch 结果的运行统计。

    从每条结果的 meta 里累加：各结果 elapsed_ms 的总和、实际 LLM 调用次数、
    各类 token 总量，以及整体缓存命中率（按 prompt_cache_hit_tokens /
    prompt_tokens 加权）。供 CLI 打印批统计用，不改变 classify_batch 的返回结构。

    注意耗时两个字段的区别（classify_batch 是并发执行的）：
      - total_elapsed_ms：各结果自身耗时(elapsed_ms)的**累加**，即"累计工作量"
        （相当于把这些任务串行跑的总时长），并发下会明显大于真实墙钟；
      - wall_clock_ms：调用方传入的整批分类**真实墙钟**耗时（在 classify_batch
        外层用 perf_counter 计时），二者之比 ≈ 并发增益（≈ max_workers）。
    """
    total_elapsed = 0.0
    llm_calls = 0
    sum_prompt = sum_completion = sum_total = sum_hit = sum_miss = 0
    for r in results:
        meta = r.get("meta") or {}
        e = meta.get("elapsed_ms")
        if e:
            total_elapsed += e
        llm = meta.get("llm")
        if isinstance(llm, dict):
            u = llm.get("usage")
            if isinstance(u, dict):
                llm_calls += 1
                sum_prompt += u.get("prompt_tokens", 0) or 0
                sum_completion += u.get("completion_tokens", 0) or 0
                sum_total += u.get("total_tokens", 0) or 0
                sum_hit += u.get("prompt_cache_hit_tokens", 0) or 0
                sum_miss += u.get("prompt_cache_miss_tokens", 0) or 0
    avg_cache_hit_rate = round(sum_hit / sum_prompt, 4) if sum_prompt else 0.0
    out = {
        "count": len(results),
        "total_elapsed_ms": round(total_elapsed, 1),
        "llm_calls": llm_calls,
        "total_prompt_tokens": sum_prompt,
        "total_completion_tokens": sum_completion,
        "total_tokens": sum_total,
        "total_cache_hit_tokens": sum_hit,
        "total_cache_miss_tokens": sum_miss,
        "avg_cache_hit_rate": avg_cache_hit_rate,
    }
    if wall_clock_ms is not None:
        out["wall_clock_ms"] = round(float(wall_clock_ms), 1)
    return out