# -*- coding: utf-8 -*-
"""速度测试执行器：校准 -> 采样 -> 汇总，全程写日志与指标入库"""
import json
import statistics
import threading
import time
import uuid

import database as db
import llm_providers as lp
from llm_providers import ProviderError, StopRequested


class TestRunner(threading.Thread):
    def __init__(self, test_id, cfg, gen):
        super().__init__(daemon=True)
        self.test_id = test_id
        self.cfg = cfg
        self.gen = gen
        self.cancel_flag = False
        self.start_wall = time.time()
        self.ratio = None
        self.samples = []
        self.last_error = None

    def request_cancel(self):
        self.cancel_flag = True

    def should_stop(self):
        return self.cancel_flag

    def log(self, level, msg):
        db.add_log(self.test_id, level, msg)

    # ───────────────────────── 主流程 ─────────────────────────

    def run(self):
        try:
            self._run()
        except StopRequested:
            self.log("WARN", "用户请求停止测试")
            db.update_status(self.test_id, "canceled",
                             summary=self._make_summary(), error="用户取消")
        except Exception as e:
            self.log("ERROR", "测试异常终止: %s" % e)
            db.update_status(self.test_id, "error",
                             summary=self._make_summary(), error=str(e))

    def _run(self):
        provider = self.cfg.get("provider", "openai")
        model = self.cfg.get("model", "")
        gen = self.gen

        # 上下文长度列表（支持手动自定义，默认 512/2048/4096/8192/16384/32768/65536/131072）
        raw_lengths = gen.get("context_lengths") or []
        if not raw_lengths:
            # 兼容旧版单值配置
            raw_lengths = [int(gen.get("prompt_tokens", 2048))]
        lengths = sorted(set(int(x) for x in raw_lengths if int(x) >= 16)) or [2048]
        n = max(1, int(gen.get("samples", 2)))                 # 每个 (长度×并发) 组合采样次数
        max_tokens = max(1, int(gen.get("max_tokens", 128)))   # 解码输出长度
        avoid_cache = bool(gen.get("avoid_cache"))
        warmup = bool(gen.get("warmup", True))                 # 测试前空转预热

        # 并发数列表（默认单流 [1]；支持 2/4 及自定义，如 [1,2,4,8]）
        raw_concs = gen.get("concurrency_levels") or []
        if not raw_concs:
            raw_concs = [int(gen.get("concurrency", 1))]
        concurrency_levels = sorted(set(int(x) for x in raw_concs if int(x) >= 1)) or [1]

        self.log("INFO", "═══ 开始速度测试 ═══")
        name = gen.get("name") or self.cfg.get("name") or ""
        if name:
            self.log("INFO", "测试名称(主题): %s" % name)
        self.log("INFO", "提供商: %s | 模型: %s" % (lp.PROVIDER_LABELS.get(provider, provider), model))
        self.log("INFO", "上下文长度: %s tokens | 生成长度: %d tokens | 并发数: %s | 每个组合采样: %d 次 | 预热: %s | 避免缓存: %s"
                 % (" / ".join(str(x) for x in lengths), max_tokens,
                    " / ".join(str(x) for x in concurrency_levels), n,
                    "开" if warmup else "关", "开" if avoid_cache else "关"))

        ratio = self._calibrate()
        self.ratio = ratio
        self.log("INFO", "校准完成: %.3f tok/字符（%.2f 字符/token）" % (ratio, 1.0 / ratio))

        run_seq = 0
        for L in lengths:
            if self.should_stop():
                raise StopRequested()
            base_prompt = self._build_prompt(L, ratio)
            self.log("INFO", "▸▸ 上下文长度 %d tokens（基准提示词构造完成）" % L)
            for C in concurrency_levels:
                if self.should_stop():
                    raise StopRequested()
                self.log("INFO", "══ 并发数 %d（同时 %d 个流）══" % (C, C))
                if warmup:
                    self._warmup(base_prompt, C)
                for i in range(1, n + 1):
                    if self.should_stop():
                        raise StopRequested()
                    run_seq += 1
                    self.log("INFO", "── [%d tok · 并发%d] 采样 %d/%d 开始 ──" % (L, C, i, n))
                    try:
                        m = self._run_sample(C, base_prompt, max_tokens, avoid_cache)
                        m["run_index"] = i
                        m["context_length"] = L
                        self.samples.append({"run_index": i, "context_length": L,
                                             "concurrency": C, "ok": True, "metrics": m})
                        db.add_run(self.test_id, run_seq, m, context_length=L)
                        self.log("METRIC", self._fmt_metric(L, C, i, n, m))
                    except StopRequested:
                        raise
                    except ProviderError as e:
                        # 单次采样失败：记录并继续后续采样，不让整个测试中断
                        self.last_error = str(e)
                        self.log("ERROR", "[%d tok · 并发%d] 采样 %d/%d 失败: %s" % (L, C, i, n, e))
                        self.samples.append({"run_index": i, "context_length": L,
                                             "concurrency": C, "ok": False, "error": str(e)})
                        db.add_run(self.test_id, run_seq, {}, str(e), context_length=L)

        summary = self._make_summary()
        ok_count = summary.get("samples_ok") or 0
        fail_count = summary.get("samples_total", 0) - ok_count
        if ok_count:
            db.update_status(self.test_id, "done", summary=summary,
                             error=("%d 次采样失败：%s" % (fail_count, self.last_error)) if fail_count else "")
            self.log("INFO", "═══ 测试完成 ═══")
            if fail_count:
                self.log("WARN", "共 %d 次采样失败（最后错误：%s）" % (fail_count, self.last_error))
        else:
            db.update_status(self.test_id, "error", summary=summary,
                             error=self.last_error or "所有采样均失败")
            self.log("ERROR", "所有采样均失败，测试标记为 error（最后错误：%s）" % (self.last_error or "未知"))
            return
        self.log("INFO", "汇总: 平均首字 %.1f ms | 平均预填充 %.1f tok/s | 平均解码 %.1f tok/s"
                 % (summary.get("avg_ttft_ms") or 0,
                    summary.get("avg_prefill_speed") or 0,
                    summary.get("avg_decode_speed") or 0))

    def _warmup(self, base_prompt, concurrency=1):
        """空转预热：不计入任何速度统计，用于避免冷启动/首次请求偏慢影响采样（按并发数预热）"""
        self.log("INFO", "预热（空转，不计速度，并发 %d）..." % concurrency)
        try:
            self._run_sample(concurrency, base_prompt, 8, False)
            self.log("INFO", "预热完成（不纳入统计）")
        except StopRequested:
            raise
        except Exception as e:
            self.log("WARN", "预热失败（继续测试）: %s" % e)

    # ───────────────────────── 并发采样 ─────────────────────────

    def _run_sample(self, concurrency, base_prompt, max_tokens, avoid_cache):
        """
        运行一个采样：concurrency 个流同时并发请求（并发=1 即单流）。
        返回聚合指标：prompt/output tokens 为 N 流之和，
        prefill/decode 速度为“整批吞吐”（tok/s），并附每流明细 streams。
        """
        from concurrent.futures import ThreadPoolExecutor

        gen_opt = {"max_tokens": max_tokens, "avoid_cache": avoid_cache}

        def worker(idx):
            prompt = self._finalize_prompt(base_prompt)  # 每流独立随机前缀，避免共享缓存
            t0 = time.time()
            try:
                m = lp.call_stream(self.cfg, prompt, gen_opt,
                                   log=lambda lv, msg: self.log(lv, msg),
                                   should_stop=self.should_stop)
                m["_wall_start"] = t0
                m["_wall_end"] = time.time()
                m["_stream_idx"] = idx
                return {"ok": True, "metrics": m}
            except StopRequested:
                raise
            except ProviderError as e:
                return {"ok": False, "error": str(e)}
            except Exception as e:
                return {"ok": False, "error": str(e)}

        with ThreadPoolExecutor(max_workers=concurrency) as ex:
            results = list(ex.map(worker, range(concurrency)))

        ok = [r["metrics"] for r in results if r.get("ok")]
        streams_detail = [self._clean_stream(m, m.get("_stream_idx", 0)) for m in ok]
        if not ok:
            errs = [r.get("error") or "未知错误" for r in results if not r.get("ok")]
            raise ProviderError("并发 %d 全部失败: %s" % (concurrency, " | ".join(errs[:3])))

        total_prompt = sum(m.get("prompt_tokens") or 0 for m in ok)
        total_output = sum(m.get("output_tokens") or 0 for m in ok)
        total_cached = sum(m.get("cached_tokens") or 0 for m in ok)
        total_pchars = sum(m.get("prompt_chars") or 0 for m in ok)
        total_ochars = sum(m.get("output_chars") or 0 for m in ok)
        batch_start = min(m["_wall_start"] for m in ok)
        # 整批首字时刻 = 任一流最早收到第一个 token 的时刻
        first_at = min(m["_wall_start"] + (m.get("ttft_ms") or 0) / 1000.0 for m in ok)
        batch_end = max(m["_wall_end"] for m in ok)
        ttft_ms = max((first_at - batch_start) * 1000.0, 0.1)
        decode_ms = max((batch_end - first_at) * 1000.0, 0.1)
        total_ms = max((batch_end - batch_start) * 1000.0, 0.1)
        prefill = (total_prompt / (ttft_ms / 1000.0)) if total_prompt else None
        decode = (total_output / (decode_ms / 1000.0)) if total_output else None

        agg = {
            "concurrency": concurrency,
            "streams_total": concurrency,
            "streams_ok": len(ok),
            "prompt_tokens": int(total_prompt),
            "output_tokens": int(total_output),
            "cached_tokens": int(total_cached),
            "prompt_chars": int(total_pchars),
            "output_chars": int(total_ochars),
            "ttft_ms": round(ttft_ms, 1),
            "decode_ms": round(decode_ms, 1),
            "total_ms": round(total_ms, 1),
            "prefill_speed": round(prefill, 1) if prefill else None,
            "decode_speed": round(decode, 1) if decode else None,
            "avg_stream_prefill": round(prefill / concurrency, 1) if prefill else None,
            "avg_stream_decode": round(decode / concurrency, 1) if decode else None,
            "streams": streams_detail,
        }
        return agg

    @staticmethod
    def _clean_stream(m, idx):
        """去掉内部 _wall 字段，保留每流可展示指标"""
        keep = {k: v for k, v in m.items() if not k.startswith("_")}
        keep["stream_idx"] = idx
        return keep

    # ───────────────────────── 工具方法 ─────────────────────────

    def _calibrate(self):
        probe = ("The quick brown fox jumps over the lazy dog. 人工智能大模型推理速度基准语料，"
                 "用于测量提示词预填充与流式解码性能。\n") * 40
        self.log("INFO", "正在校准 token/字符 比例（发送小探测请求）...")
        try:
            m = lp.call_stream(self.cfg, probe,
                               {"max_tokens": 8, "avoid_cache": False},
                               log=lambda lv, msg: self.log(lv, msg),
                               should_stop=self.should_stop)
            pt = m.get("prompt_tokens") or 0
            if pt and len(probe):
                ratio = pt / len(probe)
                self.log("INFO", "探测提示词 %d tokens / %d 字符 = %.3f tok/字符"
                         % (pt, len(probe), ratio))
                return max(ratio, 0.001)
        except StopRequested:
            raise
        except Exception as e:
            self.log("WARN", "校准失败（%s），使用默认估算 0.55 tok/字符" % e)
        return 0.55

    def _build_prompt(self, target_tokens, ratio):
        seg = ("基准语料：The quick brown fox jumps over the lazy dog. "
               "人工智能大模型推理性能测试文本，用于测量提示词预填充速度、首字延迟与流式解码吞吐。\n")
        target_chars = max(64, int(target_tokens / ratio))
        repeats = max(1, target_chars // len(seg))
        return seg * repeats

    def _finalize_prompt(self, base):
        if self.gen.get("avoid_cache"):
            return "[cache-bust %s]\n%s" % (uuid.uuid4().hex, base)
        return base

    def _fmt_metric(self, L, C, i, n, m):
        return ("[%d tok · 并发%d] 采样 %d/%d 完成 | 流 %d/%d 成功 | 提示词 %d tok | 缓存 %d tok | 首字 %s ms | 预填充 %s tok/s"
                " | 输出 %d tok | 解码 %s tok/s | 总耗时 %s ms"
                % (L, C, i, n, m.get("streams_ok") or 0, m.get("streams_total") or C,
                   m.get("prompt_tokens") or 0, m.get("cached_tokens") or 0,
                   m.get("ttft_ms"), m.get("prefill_speed"), m.get("output_tokens") or 0,
                   m.get("decode_speed"), m.get("total_ms")))

    def _make_summary(self):
        ok = [s for s in self.samples if s.get("ok")]
        base = {
            "provider": self.cfg.get("provider"),
            "model": self.cfg.get("model"),
            "gen": self.gen,
            "samples_total": len(self.samples),
            "samples_ok": len(ok),
            "concurrency_levels": sorted(set(s.get("concurrency", 1) for s in self.samples)) or [1],
            "calibration_chars_per_token": round(1 / self.ratio, 2) if self.ratio else None,
        }
        if not ok:
            return base

        def avg(ms, k):
            vals = [m[k] for m in ms if m.get(k) is not None]
            return round(statistics.mean(vals), 1) if vals else None

        # 按上下文长度分组汇总
        by_length = {}
        for L in sorted(set(s["context_length"] for s in ok)):
            group = [s["metrics"] for s in ok if s["context_length"] == L]
            by_length[L] = {
                "samples_total": sum(1 for s in self.samples if s["context_length"] == L),
                "samples_ok": len(group),
                "avg_ttft_ms": avg(group, "ttft_ms"),
                "avg_prefill_speed": avg(group, "prefill_speed"),
                "avg_decode_speed": avg(group, "decode_speed"),
                "avg_prompt_tokens": avg(group, "prompt_tokens"),
                "avg_output_tokens": avg(group, "output_tokens"),
                "avg_total_ms": avg(group, "total_ms"),
            }

        # 按并发数分组汇总（多测试结果并排对比的核心数据）
        by_concurrency = {}
        for C in sorted(set(s.get("concurrency", 1) for s in ok)):
            group = [s["metrics"] for s in ok if s.get("concurrency", 1) == C]
            by_concurrency[C] = {
                "samples_total": sum(1 for s in self.samples if s.get("concurrency", 1) == C),
                "samples_ok": len(group),
                "avg_ttft_ms": avg(group, "ttft_ms"),
                "avg_prefill_speed": avg(group, "prefill_speed"),
                "avg_decode_speed": avg(group, "decode_speed"),
                "avg_stream_decode": avg(group, "avg_stream_decode"),
                "avg_prompt_tokens": avg(group, "prompt_tokens"),
                "avg_output_tokens": avg(group, "output_tokens"),
                "avg_total_ms": avg(group, "total_ms"),
            }

        # 长度 × 并发 全网格（详情/Excel 用）
        by_length_concurrency = {}
        for L in sorted(set(s["context_length"] for s in ok)):
            grid = {}
            for C in sorted(set(s.get("concurrency", 1) for s in ok)):
                group = [s["metrics"] for s in ok
                         if s["context_length"] == L and s.get("concurrency", 1) == C]
                grid[C] = {
                    "samples_total": sum(1 for s in self.samples
                                          if s["context_length"] == L and s.get("concurrency", 1) == C),
                    "samples_ok": len(group),
                    "avg_ttft_ms": avg(group, "ttft_ms"),
                    "avg_prefill_speed": avg(group, "prefill_speed"),
                    "avg_decode_speed": avg(group, "decode_speed"),
                    "avg_prompt_tokens": avg(group, "prompt_tokens"),
                    "avg_output_tokens": avg(group, "output_tokens"),
                    "avg_total_ms": avg(group, "total_ms"),
                }
            by_length_concurrency[L] = grid

        okm = [s["metrics"] for s in ok]

        def mn(k):
            vals = [m[k] for m in okm if m.get(k) is not None]
            return round(min(vals), 1) if vals else None

        def mx(k):
            vals = [m[k] for m in okm if m.get(k) is not None]
            return round(max(vals), 1) if vals else None

        summary = dict(base)
        summary.update({
            "by_length": by_length,
            "by_concurrency": by_concurrency,
            "by_length_concurrency": by_length_concurrency,
            "avg_ttft_ms": avg(okm, "ttft_ms"),
            "min_ttft_ms": mn("ttft_ms"),
            "max_ttft_ms": mx("ttft_ms"),
            "avg_prefill_speed": avg(okm, "prefill_speed"),
            "min_prefill_speed": mn("prefill_speed"),
            "max_prefill_speed": mx("prefill_speed"),
            "avg_decode_speed": avg(okm, "decode_speed"),
            "min_decode_speed": mn("decode_speed"),
            "max_decode_speed": mx("decode_speed"),
            "avg_prompt_tokens": avg(okm, "prompt_tokens"),
            "avg_output_tokens": avg(okm, "output_tokens"),
            "avg_cached_tokens": avg(okm, "cached_tokens"),
            "avg_total_ms": avg(okm, "total_ms"),
            "min_total_ms": mn("total_ms"),
            "max_total_ms": mx("total_ms"),
            "best_ttft_ms": mn("ttft_ms"),
        })
        return summary