# -*- coding: utf-8 -*-
"""LLM 速度测试台 - Flask 主应用"""
import io
import json

from flask import Flask, jsonify, request, send_file, send_from_directory

import config
import database as db
from llm_providers import DEFAULT_URLS, ProviderError, call_stream
from tester import TestRunner

app = Flask(__name__, static_folder="static", static_url_path="")
app.json.ensure_ascii = False

db.init_db()

RUNNERS = {}  # test_id -> TestRunner


@app.route("/")
def index():
    return send_from_directory(app.static_folder, "index.html")


@app.route("/api/health")
def health():
    running = [tid for tid, r in RUNNERS.items() if r.is_alive()]
    return jsonify({"ok": True, "port": config.PORT, "running_tests": running})


def _fill_defaults(cfg):
    p = cfg.get("provider", "openai")
    if not cfg.get("base_url"):
        cfg["base_url"] = DEFAULT_URLS.get(p, "")
    if cfg.get("temperature") is None:
        cfg["temperature"] = 0.7
    return cfg


# ───────────────────────── 提供商配置 ─────────────────────────

@app.route("/api/configs", methods=["GET"])
def list_configs():
    return jsonify(db.list_configs())


@app.route("/api/configs", methods=["POST"])
def add_config():
    cfg = request.get_json(force=True) or {}
    if not cfg.get("name"):
        return jsonify({"ok": False, "error": "请填写配置名称"}), 400
    cid = db.add_config(cfg)
    return jsonify({"ok": True, "id": cid})


@app.route("/api/configs/<int:cid>", methods=["GET"])
def get_one_config(cid):
    c = db.get_config(cid)
    if not c:
        return jsonify({"ok": False, "error": "配置不存在"}), 404
    return jsonify(c)


@app.route("/api/configs/<int:cid>", methods=["PUT"])
def update_config(cid):
    cfg = request.get_json(force=True) or {}
    old = db.get_config(cid)
    if not old:
        return jsonify({"ok": False, "error": "配置不存在"}), 404
    db.update_config(cid, cfg)
    return jsonify({"ok": True, "id": cid})


@app.route("/api/configs/<int:cid>", methods=["DELETE"])
def del_config(cid):
    db.delete_config(cid)
    return jsonify({"ok": True})


@app.route("/api/configs/test", methods=["POST"])
def test_config():
    cfg = _fill_defaults(request.get_json(force=True) or {})
    if not cfg.get("api_key"):
        return jsonify({"ok": False, "error": "请填写 API Key"}), 400
    try:
        # 连接测试：只要流式请求成功返回（哪怕正文为空/只有思维链）都算连通
        m = call_stream(cfg, "你好，请简要回答：1+1=？",
                        {"max_tokens": 32, "avoid_cache": False})
        note = ""
        if not (m.get("output_tokens") or m.get("output_chars")):
            note = "（连接正常，但本次未返回正文内容，可能为推理型模型）"
        return jsonify({"ok": True, "total_ms": m["total_ms"], "metrics": m, "note": note})
    except ProviderError as e:
        return jsonify({"ok": False, "error": str(e)})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})


# ───────────────────────── 测试 ─────────────────────────

@app.route("/api/tests", methods=["POST"])
def start_test():
    body = request.get_json(force=True) or {}
    cfg = _fill_defaults(body.get("config") or {})
    gen = body.get("gen") or {}
    if not cfg.get("api_key"):
        return jsonify({"ok": False, "error": "请填写 API Key"}), 400
    if not cfg.get("model"):
        return jsonify({"ok": False, "error": "请填写模型名称"}), 400
    tid = db.create_test(cfg, gen)
    runner = TestRunner(tid, cfg, gen)
    RUNNERS[tid] = runner
    runner.start()
    return jsonify({"ok": True, "id": tid})


@app.route("/api/tests", methods=["GET"])
def list_tests():
    try:
        limit = int(request.args.get("limit", 100))
    except ValueError:
        limit = 100
    return jsonify(db.list_tests(max(1, min(limit, 1000))))


@app.route("/api/tests/<int:tid>", methods=["GET"])
def get_test(tid):
    t = db.get_test(tid)
    if not t:
        return jsonify({"ok": False, "error": "测试不存在"}), 404
    t["runs"] = db.get_runs(tid)
    t["logs"] = db.get_logs(tid)
    _mask_cfg(t.get("config"))
    return jsonify(t)


@app.route("/api/tests/<int:tid>/logs", methods=["GET"])
def get_logs(tid):
    after = int(request.args.get("after", 0))
    data = db.get_logs_after(tid, after)
    if data is None:
        return jsonify({"ok": False, "error": "测试不存在"}), 404
    return jsonify(data)


@app.route("/api/tests/<int:tid>/cancel", methods=["POST"])
def cancel_test(tid):
    r = RUNNERS.get(tid)
    if r and r.is_alive():
        r.request_cancel()
        return jsonify({"ok": True, "msg": "正在停止..."})
    return jsonify({"ok": False, "msg": "测试未在运行"})


@app.route("/api/tests/<int:tid>", methods=["DELETE"])
def del_test(tid):
    db.delete_test(tid)
    RUNNERS.pop(tid, None)
    return jsonify({"ok": True})


@app.route("/api/tests/<int:tid>/export.json")
def export_json(tid):
    t = db.get_test(tid)
    if not t:
        return jsonify({"ok": False, "error": "测试不存在"}), 404
    t["runs"] = db.get_runs(tid)
    t["logs"] = db.get_logs(tid)
    _mask_cfg(t.get("config"))
    return jsonify(t)


def _mask_cfg(cfg):
    """对外隐藏 API Key，仅保留前几位便于识别"""
    if cfg and cfg.get("api_key"):
        k = cfg["api_key"]
        cfg["api_key"] = k[:4] + "****" if len(k) > 6 else "****"
    return cfg


# ───────────────────────── Excel 导出 ─────────────────────────

@app.route("/api/tests/<int:tid>/export.xlsx")
def export_xlsx(tid):
    t = db.get_test(tid)
    if not t:
        return jsonify({"ok": False, "error": "测试不存在"}), 404
    t["runs"] = db.get_runs(tid)
    t["logs"] = db.get_logs(tid)
    try:
        data = _build_xlsx(t)
    except Exception as e:
        return jsonify({"ok": False, "error": "导出失败: %s" % e}), 500
    return send_file(data, as_attachment=True,
                     download_name="llm_speed_test_%d.xlsx" % tid,
                     mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")


def _build_xlsx(t):
    from openpyxl import Workbook
    from openpyxl.styles import Alignment, Font, PatternFill

    s = t.get("summary") or {}
    g = t.get("gen") or {}
    cfg = t.get("config") or {}
    by_length = s.get("by_length") or {}
    runs = t.get("runs") or []
    logs = t.get("logs") or []

    wb = Workbook()
    head_fill = PatternFill("solid", fgColor="2A3550")
    head_font = Font(color="FFFFFF", bold=True)
    title_font = Font(bold=True, size=12)

    def style_header(ws, row, ncol):
        for c in range(1, ncol + 1):
            cell = ws.cell(row=row, column=c)
            cell.fill = head_fill
            cell.font = head_font
            cell.alignment = Alignment(horizontal="center", vertical="center")

    # ── Sheet1 汇总 ──
    ws = wb.active
    ws.title = "汇总"
    ws.append(["LLM 速度测试报告"])
    ws.cell(1, 1).font = Font(bold=True, size=14)
    info = [
        ["测试编号", "#%d" % t["id"]],
        ["测试名称", t.get("name") or "（未命名）"],
        ["创建时间", t.get("created_at", "")],
        ["状态", t.get("status", "")],
        ["提供商", t.get("provider", "")],
        ["模型", t.get("model", "")],
        ["Base URL", cfg.get("base_url") or "（默认）"],
        ["上下文长度列表", " / ".join(str(x) for x in (g.get("context_lengths") or []))],
        ["生成长度(max tokens)", g.get("max_tokens", 128)],
        ["每个长度采样次数", g.get("samples", 2)],
        ["预热(空转)", "开" if g.get("warmup", True) else "关"],
        ["避免缓存", "开" if g.get("avoid_cache") else "关"],
        ["采样(成功/总数)", "%s / %s" % (s.get("samples_ok"), s.get("samples_total"))],
        ["校准 字符/token", s.get("calibration_chars_per_token") or "—"],
        ["错误信息", t.get("error") or ""],
    ]
    for row in info:
        ws.append(row)
    ws.cell(15, 1).font = title_font
    r0 = len(info) + 2
    overall = [
        ["首字延迟(ms)", s.get("avg_ttft_ms"), s.get("max_ttft_ms"), s.get("min_ttft_ms")],
        ["预填充速度(tok/s)", s.get("avg_prefill_speed"), s.get("max_prefill_speed"), s.get("min_prefill_speed")],
        ["解码速度(tok/s)", s.get("avg_decode_speed"), s.get("max_decode_speed"), s.get("min_decode_speed")],
        ["提示词(tok)", s.get("avg_prompt_tokens"), None, None],
        ["输出(tok)", s.get("avg_output_tokens"), None, None],
        ["总耗时(ms)", s.get("avg_total_ms"), s.get("max_total_ms"), s.get("min_total_ms")],
    ]
    ws.cell(r0, 1, "整体统计指标（平均 / 最大 / 最小）").font = title_font
    for j, c in enumerate(["指标", "平均", "最大", "最小"], start=1):
        ws.cell(row=r0 + 1, column=j, value=c)
    style_header(ws, r0 + 1, 4)
    for i, row in enumerate(overall, start=r0 + 2):
        for j, v in enumerate(row, start=1):
            ws.cell(row=i, column=j, value=v)

    # 按上下文长度分组
    r1 = r0 + len(overall) + 3
    ws.cell(r1, 1, "按上下文长度分组").font = title_font
    cols = ["上下文长度(tok)", "采样(成功/总数)", "首字ms", "预填充tok/s", "解码tok/s", "提示词tok", "输出tok", "总耗时ms"]
    ws.append([])
    for j, c in enumerate(cols, start=1):
        ws.cell(row=r1 + 1, column=j, value=c)
    style_header(ws, r1 + 1, len(cols))
    if by_length:
        rr = r1 + 2
        for L in sorted(int(k) for k in by_length):
            bl = by_length[str(L)] if str(L) in by_length else by_length[L]
            ws.cell(row=rr, column=1, value=L)
            ws.cell(row=rr, column=2, value="%s / %s" % (bl.get("samples_ok"), bl.get("samples_total")))
            ws.cell(row=rr, column=3, value=bl.get("avg_ttft_ms"))
            ws.cell(row=rr, column=4, value=bl.get("avg_prefill_speed"))
            ws.cell(row=rr, column=5, value=bl.get("avg_decode_speed"))
            ws.cell(row=rr, column=6, value=bl.get("avg_prompt_tokens"))
            ws.cell(row=rr, column=7, value=bl.get("avg_output_tokens"))
            ws.cell(row=rr, column=8, value=bl.get("avg_total_ms"))
            rr += 1
    else:
        ws.cell(row=r1 + 2, column=1, value="（无成功采样数据）")
    for col, w in zip("ABCDEFGH", [22, 20, 12, 14, 14, 12, 12, 14]):
        ws.column_dimensions[col].width = w

    # ── Sheet2 采样明细 ──
    ws2 = wb.create_sheet("采样明细")
    h2 = ["序号", "上下文长度tok", "提示词tok", "缓存tok", "首字ms", "预填充tok/s",
          "输出tok", "解码tok/s", "总耗时ms", "备注"]
    ws2.append(h2)
    style_header(ws2, 1, len(h2))
    for i, r in enumerate(runs, start=1):
        m = r.get("metrics") or {}
        ws2.append([
            i,
            r.get("context_length") or m.get("context_length") or "",
            m.get("prompt_tokens") or "",
            m.get("cached_tokens") if m.get("cached_tokens") else "",
            m.get("ttft_ms"),
            m.get("prefill_speed"),
            m.get("output_tokens"),
            m.get("decode_speed"),
            m.get("total_ms"),
            r.get("error") or "OK",
        ])
    for col, w in zip("ABCDEFGHIJ", [8, 14, 12, 10, 12, 14, 12, 14, 12, 30]):
        ws2.column_dimensions[col].width = w

    # ── Sheet3 日志 ──
    ws3 = wb.create_sheet("日志")
    ws3.append(["相对时间(s)", "级别", "内容"])
    style_header(ws3, 1, 3)
    for l in logs:
        ws3.append([l.get("rel", 0), l.get("level", ""), l.get("msg", "")])
    for col, w in zip("ABC", [14, 10, 90]):
        ws3.column_dimensions[col].width = w

    bio = io.BytesIO()
    wb.save(bio)
    bio.seek(0)
    return bio


if __name__ == "__main__":
    app.run(host=config.HOST, port=config.PORT, threaded=True, debug=False)