# -*- coding: utf-8 -*-
"""SQLite 存储：提供商配置 / 测试记录 / 每次采样指标 / 全量日志"""
import os
import json
import math
import sqlite3
import threading
import time

import config

_lock = threading.RLock()  # RLock：允许 get_logs_after 内嵌套调用 get_last_run

SCHEMA = """
CREATE TABLE IF NOT EXISTS configs(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  name TEXT DEFAULT '',
  provider TEXT DEFAULT 'openai',
  base_url TEXT DEFAULT '',
  api_key TEXT DEFAULT '',
  model TEXT DEFAULT '',
  temperature REAL DEFAULT 0.7,
  created_at TEXT DEFAULT (datetime('now','localtime'))
);

CREATE TABLE IF NOT EXISTS tests(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  created_at TEXT DEFAULT (datetime('now','localtime')),
  status TEXT DEFAULT 'running',
  provider TEXT DEFAULT '',
  model TEXT DEFAULT '',
  name TEXT DEFAULT '',
  config_json TEXT DEFAULT '{}',
  gen_cfg_json TEXT DEFAULT '{}',
  summary_json TEXT DEFAULT '{}',
  error TEXT DEFAULT '',
  started_at REAL DEFAULT 0
);

CREATE TABLE IF NOT EXISTS test_runs(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  test_id INTEGER NOT NULL,
  run_index INTEGER DEFAULT 0,
  context_length INTEGER DEFAULT 0,
  metrics_json TEXT DEFAULT '{}',
  error TEXT DEFAULT ''
);

CREATE TABLE IF NOT EXISTS logs(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  test_id INTEGER NOT NULL,
  level TEXT DEFAULT 'INFO',
  msg TEXT DEFAULT '',
  rel REAL DEFAULT 0,
  ts TEXT DEFAULT (datetime('now','localtime'))
);

CREATE INDEX IF NOT EXISTS idx_runs_test ON test_runs(test_id);
CREATE INDEX IF NOT EXISTS idx_logs_test ON logs(test_id);
"""


def _connect():
    os.makedirs(config.DATA_DIR, exist_ok=True)
    conn = sqlite3.connect(config.DB_PATH, check_same_thread=False, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    return conn


def _migrate(conn):
    """老库升级：test_runs 补 context_length、tests 补 name"""
    cur = conn.execute("PRAGMA table_info(test_runs)")
    cols = [r[1] for r in cur.fetchall()]
    if "context_length" not in cols:
        conn.execute("ALTER TABLE test_runs ADD COLUMN context_length INTEGER DEFAULT 0")
    cur = conn.execute("PRAGMA table_info(tests)")
    tcols = [r[1] for r in cur.fetchall()]
    if "name" not in tcols:
        conn.execute("ALTER TABLE tests ADD COLUMN name TEXT DEFAULT ''")


def init_db():
    with _lock:
        conn = _connect()
        try:
            conn.executescript(SCHEMA)
            _migrate(conn)
            conn.commit()
        finally:
            conn.close()


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

def add_config(cfg: dict) -> int:
    with _lock:
        conn = _connect()
        try:
            cur = conn.execute(
                "INSERT INTO configs(name,provider,base_url,api_key,model,temperature) VALUES(?,?,?,?,?,?)",
                (cfg.get("name", ""), cfg.get("provider", "openai"),
                 cfg.get("base_url", ""), cfg.get("api_key", ""),
                 cfg.get("model", ""), float(cfg.get("temperature", 0.7))))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def list_configs():
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute("SELECT id,name,provider,base_url,model,temperature,"
                                "CASE WHEN api_key<>'' THEN 1 ELSE 0 END AS has_key "
                                "FROM configs ORDER BY id DESC").fetchall()
            return [dict(r) for r in rows]
        finally:
            conn.close()


def get_config(cid: int):
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT * FROM configs WHERE id=?", (cid,)).fetchone()
            return dict(r) if r else None
        finally:
            conn.close()


def delete_config(cid: int):
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM configs WHERE id=?", (cid,))
            conn.commit()
        finally:
            conn.close()


def update_config(cid: int, cfg: dict):
    """局部更新：只更新请求里出现的字段"""
    allow = {"name", "provider", "base_url", "api_key", "model", "temperature"}
    fields = {k: v for k, v in cfg.items() if k in allow and v is not None}
    if not fields:
        return
    sets = ",".join("%s=?" % k for k in fields)
    vals = list(fields.values()) + [cid]
    with _lock:
        conn = _connect()
        try:
            conn.execute("UPDATE configs SET %s WHERE id=?" % sets, vals)
            conn.commit()
        finally:
            conn.close()


# ───────────────────────── 测试记录 ─────────────────────────

def create_test(cfg: dict, gen: dict) -> int:
    with _lock:
        conn = _connect()
        try:
            cur = conn.execute(
                "INSERT INTO tests(status,provider,model,name,config_json,gen_cfg_json,started_at) "
                "VALUES('running',?,?,?,?,?,?)",
                (cfg.get("provider", "openai"), cfg.get("model", ""),
                 gen.get("name") or cfg.get("name") or "",
                 json.dumps(cfg, ensure_ascii=False), json.dumps(gen, ensure_ascii=False),
                 time.time()))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def update_status(tid: int, status: str, summary: dict = None, error: str = ""):
    with _lock:
        conn = _connect()
        try:
            conn.execute("UPDATE tests SET status=?, summary_json=?, error=? WHERE id=?",
                         (status,
                          json.dumps(summary or {}, ensure_ascii=False),
                          error or "", tid))
            conn.commit()
        finally:
            conn.close()


def get_test(tid: int):
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT * FROM tests WHERE id=?", (tid,)).fetchone()
            if not r:
                return None
            d = dict(r)
            d["config"] = json.loads(d.pop("config_json") or "{}")
            d["gen"] = json.loads(d.pop("gen_cfg_json") or "{}")
            d["summary"] = json.loads(d.pop("summary_json") or "{}")
            return d
        finally:
            conn.close()


# 历史列表排序字段 → 取值函数（数值字段取 summary 聚合值，None 恒排最后）
SORT_KEYS = {
    "id": lambda t: t["id"],
    "created_at": lambda t: t["created_at"],
    "name": lambda t: (t.get("name") or "").lower(),
    "provider": lambda t: (t.get("provider") or "").lower(),
    "model": lambda t: (t.get("model") or "").lower(),
    "status": lambda t: t.get("status") or "",
    "samples": lambda t: (t.get("summary") or {}).get("samples_ok"),
    "ttft": lambda t: (t.get("summary") or {}).get("avg_ttft_ms"),
    "prefill": lambda t: (t.get("summary") or {}).get("avg_prefill_speed"),
    "decode": lambda t: (t.get("summary") or {}).get("avg_decode_speed"),
    "total_ms": lambda t: (t.get("summary") or {}).get("avg_total_ms"),
}


def query_tests(q="", status="", provider="", sort="id", order="desc",
                page=1, page_size=20):
    """历史记录：关键词搜索 + 状态/提供商筛选 + 字段排序 + 分页。
    返回 {"items": [...], "total": N, "page": p, "page_size": s, "total_pages": n}"""
    with _lock:
        conn = _connect()
        try:
            where, args = [], []
            if q:
                like = "%" + q + "%"
                where.append("(id LIKE ? OR name LIKE ? OR model LIKE ? OR provider LIKE ?"
                             " OR status LIKE ? OR created_at LIKE ?)")
                args += [like] * 6
            if status:
                sts = [s.strip() for s in status.split(",") if s.strip()]
                if sts:
                    where.append("status IN (%s)" % ",".join("?" * len(sts)))
                    args += sts
            if provider:
                where.append("provider = ?")
                args.append(provider)
            wsql = ("WHERE " + " AND ".join(where)) if where else ""

            rows = conn.execute(
                "SELECT id,created_at,status,provider,model,name,summary_json,gen_cfg_json,error "
                "FROM tests " + wsql + " ORDER BY id DESC", args).fetchall()
            out = []
            for r in rows:
                d = dict(r)
                d["summary"] = json.loads(d.pop("summary_json") or "{}")
                d["gen"] = json.loads(d.pop("gen_cfg_json") or "{}")
                out.append(d)

            # Python 侧排序：None 值恒排最后（无论升降序）
            keyfn = SORT_KEYS.get(sort) or SORT_KEYS["id"]
            present = [t for t in out if keyfn(t) is not None]
            absent = [t for t in out if keyfn(t) is None]
            present.sort(key=keyfn, reverse=(order == "desc"))
            out = present + absent

            total = len(out)
            total_pages = max(1, math.ceil(total / page_size))
            page = min(max(1, page), total_pages)
            start = (page - 1) * page_size
            return {
                "items": out[start:start + page_size],
                "total": total,
                "page": page,
                "page_size": page_size,
                "total_pages": total_pages,
            }
        finally:
            conn.close()


def clear_tests():
    """一键清空全部历史（测试记录 + 采样指标 + 日志）"""
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM test_runs")
            conn.execute("DELETE FROM logs")
            conn.execute("DELETE FROM tests")
            conn.commit()
        finally:
            conn.close()


def delete_test(tid: int):
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM tests WHERE id=?", (tid,))
            conn.execute("DELETE FROM test_runs WHERE test_id=?", (tid,))
            conn.execute("DELETE FROM logs WHERE test_id=?", (tid,))
            conn.commit()
        finally:
            conn.close()


# ───────────────────────── 采样指标 ─────────────────────────

def add_run(tid: int, run_index: int, metrics: dict, error: str = "", context_length: int = 0):
    with _lock:
        conn = _connect()
        try:
            conn.execute(
                "INSERT INTO test_runs(test_id,run_index,context_length,metrics_json,error) VALUES(?,?,?,?,?)",
                (tid, run_index, context_length, json.dumps(metrics, ensure_ascii=False), error))
            conn.commit()
        finally:
            conn.close()


def get_runs(tid: int):
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT run_index,context_length,metrics_json,error FROM test_runs "
                "WHERE test_id=? ORDER BY run_index", (tid,)).fetchall()
            out = []
            for r in rows:
                d = dict(r)
                d["metrics"] = json.loads(d.pop("metrics_json") or "{}")
                out.append(d)
            return out
        finally:
            conn.close()


def get_last_run(tid: int):
    with _lock:
        conn = _connect()
        try:
            r = conn.execute(
                "SELECT metrics_json FROM test_runs WHERE test_id=? "
                "ORDER BY run_index DESC LIMIT 1", (tid,)).fetchone()
            return json.loads(r["metrics_json"] or "{}") if r else None
        finally:
            conn.close()


# ───────────────────────── 日志 ─────────────────────────

def add_log(tid: int, level: str, msg: str, rel: float = None):
    with _lock:
        conn = _connect()
        try:
            if rel is None:
                r = conn.execute("SELECT started_at FROM tests WHERE id=?", (tid,)).fetchone()
                rel = (time.time() - (r["started_at"] or time.time())) if r else 0.0
            conn.execute("INSERT INTO logs(test_id,level,msg,rel) VALUES(?,?,?,?)",
                         (tid, level, msg, round(rel, 3)))
            conn.commit()
        finally:
            conn.close()


def get_logs_after(tid: int, after_id: int = 0):
    with _lock:
        conn = _connect()
        try:
            t = conn.execute("SELECT status,summary_json,error FROM tests WHERE id=?", (tid,)).fetchone()
            if not t:
                return None
            rows = conn.execute(
                "SELECT id,level,msg,rel FROM logs WHERE test_id=? AND id>? ORDER BY id",
                (tid, after_id)).fetchall()
            logs = [dict(r) for r in rows]
            last = logs[-1]["id"] if logs else after_id
            return {
                "status": t["status"],
                "error": t["error"],
                "summary": json.loads(t["summary_json"] or "{}"),
                "last_run": get_last_run(tid),
                "logs": logs,
                "after": last,
            }
        finally:
            conn.close()


def get_logs(tid: int):
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT id,level,msg,rel,ts FROM logs WHERE test_id=? ORDER BY id", (tid,)).fetchall()
            return [dict(r) for r in rows]
        finally:
            conn.close()