# -*- coding: utf-8 -*-
"""SQLite 存储：提供商配置 / 测试记录 / 每次采样指标 / 全量日志"""
import os
import json
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 '',
  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 列"""
    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")


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 create_test(cfg: dict, gen: dict) -> int:
    with _lock:
        conn = _connect()
        try:
            cur = conn.execute(
                "INSERT INTO tests(status,provider,model,config_json,gen_cfg_json,started_at) "
                "VALUES('running',?,?,?,?,?)",
                (cfg.get("provider", "openai"), cfg.get("model", ""),
                 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()


def list_tests(limit=100):
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT id,created_at,status,provider,model,summary_json,error "
                "FROM tests ORDER BY id DESC LIMIT ?", (limit,)).fetchall()
            out = []
            for r in rows:
                d = dict(r)
                d["summary"] = json.loads(d.pop("summary_json") or "{}")
                out.append(d)
            return out
        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()