# -*- coding: utf-8 -*-
"""
数据库层：SQLite + WAL，轻量直连封装
"""
import sqlite3
import json
import time
from config import DB_PATH

SCHEMA = """
CREATE TABLE IF NOT EXISTS projects (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    objective TEXT DEFAULT '',
    acceptance_criteria TEXT DEFAULT '',
    status TEXT DEFAULT 'active',          -- planning/active/done/archived
    budget_limit REAL DEFAULT 0,           -- 项目预算上限（元），0=不限
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS workers (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    provider TEXT NOT NULL,
    model TEXT NOT NULL,
    base_url TEXT DEFAULT '',
    api_key TEXT DEFAULT '',               -- 留空则用供应商全局 key
    system_prompt TEXT DEFAULT '',
    temperature REAL DEFAULT 0.7,
    max_tokens INTEGER DEFAULT 2000,
    task_cost_limit REAL DEFAULT 0,        -- 单任务成本上限（元），0=不限
    monthly_cost_limit REAL DEFAULT 0,     -- 月度成本上限（元），0=不限
    status TEXT DEFAULT 'enabled',         -- enabled/disabled
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS tasks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    project_id INTEGER NOT NULL,
    worker_id INTEGER,                     -- NULL = 自动路由
    title TEXT NOT NULL,
    description TEXT DEFAULT '',
    status TEXT DEFAULT 'todo',            -- todo/running/review/done/rejected/failed/cancelled
    priority TEXT DEFAULT 'medium',        -- high/medium/low
    review_required INTEGER DEFAULT 1,     -- 完成后是否需要人工审核
    output_text TEXT DEFAULT '',
    output_version INTEGER DEFAULT 0,
    rejection_count INTEGER DEFAULT 0,
    error TEXT DEFAULT '',
    deadline TEXT DEFAULT '',
    created_at INTEGER,
    updated_at INTEGER,
    started_at INTEGER,
    finished_at INTEGER
);

CREATE TABLE IF NOT EXISTS task_logs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id INTEGER NOT NULL,
    level TEXT DEFAULT 'info',             -- info/success/warn/error
    message TEXT DEFAULT '',
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS cost_records (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id INTEGER,
    project_id INTEGER,
    worker_id INTEGER,
    provider TEXT DEFAULT '',
    model TEXT DEFAULT '',
    prompt_tokens INTEGER DEFAULT 0,
    completion_tokens INTEGER DEFAULT 0,
    total_tokens INTEGER DEFAULT 0,
    cost REAL DEFAULT 0,
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS documents (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    project_id INTEGER NOT NULL,
    name TEXT NOT NULL,
    content TEXT DEFAULT '',
    source TEXT DEFAULT 'manual',          -- manual/file/url
    chunk_size INTEGER DEFAULT 0,
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS doc_chunks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    document_id INTEGER NOT NULL,
    idx INTEGER DEFAULT 0,
    content TEXT DEFAULT '',
    tokens INTEGER DEFAULT 0
);

CREATE TABLE IF NOT EXISTS alerts (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    type TEXT DEFAULT 'system',            -- budget/task_failed/worker/limit/notify/plan
    level TEXT DEFAULT 'info',             -- info/warn/critical
    title TEXT DEFAULT '',
    detail TEXT DEFAULT '',
    read INTEGER DEFAULT 0,
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS api_tokens (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    token TEXT NOT NULL UNIQUE,
    created_at INTEGER,
    last_used_at INTEGER
);

CREATE TABLE IF NOT EXISTS notify_channels (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    type TEXT NOT NULL,                    -- feishu/wecom/email
    webhook TEXT DEFAULT '',
    email TEXT DEFAULT '',
    events TEXT DEFAULT '[]',              -- JSON: task_review/task_done/task_failed/budget_alert/worker_alert
    enabled INTEGER DEFAULT 1,
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS settings (
    key TEXT PRIMARY KEY,
    value TEXT DEFAULT ''
);

CREATE INDEX IF NOT EXISTS idx_tasks_project ON tasks(project_id);

-- ===================================================================
-- V2 表结构：多 Agent 协作 / 自动评估 / 模板市场 / 企业版
-- ===================================================================
CREATE TABLE IF NOT EXISTS agent_runs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    mode TEXT NOT NULL,                -- supervisor / review / debate
    title TEXT DEFAULT '',
    topic TEXT DEFAULT '',             -- 输入主题 / 任务
    context TEXT DEFAULT '',           -- 附加上下文（知识库/约束）
    worker_ids TEXT DEFAULT '[]',      -- JSON: 参与协作的 worker id 列表
    params TEXT DEFAULT '{}',          -- JSON: rounds/阈值/立场等
    status TEXT DEFAULT 'running',     -- running/done/failed/cancelled
    result TEXT DEFAULT '',            -- 最终产出
    summary TEXT DEFAULT '',           -- 过程摘要（评审意见/共识等）
    error TEXT DEFAULT '',
    total_tokens INTEGER DEFAULT 0,
    cost REAL DEFAULT 0,
    created_at INTEGER,
    finished_at INTEGER
);

CREATE TABLE IF NOT EXISTS agent_steps (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    run_id INTEGER NOT NULL,
    role TEXT DEFAULT '',              -- supervisor/worker/reviewer/judge/debater
    worker_id INTEGER,
    seq INTEGER DEFAULT 0,
    stage TEXT DEFAULT '',             -- plan/delegate/produce/critique/revise/synthesize/verdict
    content TEXT DEFAULT '',
    tokens INTEGER DEFAULT 0,
    cost REAL DEFAULT 0,
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS eval_datasets (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    rubric TEXT DEFAULT '',            -- 评分标准（LLM-as-judge）
    tags TEXT DEFAULT '[]',
    is_builtin INTEGER DEFAULT 0,
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS eval_cases (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    dataset_id INTEGER NOT NULL,
    input TEXT DEFAULT '',
    expected TEXT DEFAULT '',
    tags TEXT DEFAULT '[]',
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS eval_runs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    dataset_id INTEGER NOT NULL,
    worker_id INTEGER NOT NULL,
    status TEXT DEFAULT 'running',     -- running/done/failed/cancelled
    score REAL DEFAULT 0,              -- 平均分 0-100
    total_tokens INTEGER DEFAULT 0,
    cost REAL DEFAULT 0,
    cases_total INTEGER DEFAULT 0,
    cases_done INTEGER DEFAULT 0,
    created_at INTEGER,
    finished_at INTEGER
);

CREATE TABLE IF NOT EXISTS eval_results (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    run_id INTEGER NOT NULL,
    case_id INTEGER NOT NULL,
    worker_id INTEGER,
    output TEXT DEFAULT '',
    score REAL DEFAULT 0,
    judgment TEXT DEFAULT '',
    latency_ms INTEGER DEFAULT 0,
    cost REAL DEFAULT 0,
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS templates (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    type TEXT NOT NULL,                -- task / project / team
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    content TEXT DEFAULT '{}',         -- JSON
    tags TEXT DEFAULT '[]',
    author TEXT DEFAULT 'system',
    is_builtin INTEGER DEFAULT 0,
    usage_count INTEGER DEFAULT 0,
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS users (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT NOT NULL UNIQUE,
    password_hash TEXT DEFAULT '',
    display_name TEXT DEFAULT '',
    role TEXT DEFAULT 'member',        -- admin / member / auditor
    source TEXT DEFAULT 'local',       -- local / oidc / ldap
    status TEXT DEFAULT 'active',      -- active / disabled
    last_login_at INTEGER,
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS audit_logs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    actor TEXT DEFAULT '',             -- 用户名 / token 名 / system
    action TEXT DEFAULT '',            -- 如 task.create / worker.update / agent.run
    target TEXT DEFAULT '',
    detail TEXT DEFAULT '',
    ip TEXT DEFAULT '',
    user_agent TEXT DEFAULT '',
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS enterprise_settings (
    key TEXT PRIMARY KEY,
    value TEXT DEFAULT ''
);

CREATE INDEX IF NOT EXISTS idx_agent_steps_run ON agent_steps(run_id);
CREATE INDEX IF NOT EXISTS idx_agent_runs_status ON agent_runs(status);
CREATE INDEX IF NOT EXISTS idx_eval_cases_ds ON eval_cases(dataset_id);
CREATE INDEX IF NOT EXISTS idx_eval_runs_ds ON eval_runs(dataset_id);
CREATE INDEX IF NOT EXISTS idx_eval_results_run ON eval_results(run_id);
CREATE INDEX IF NOT EXISTS idx_audit_time ON audit_logs(created_at);
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_logs_task ON task_logs(task_id);
CREATE INDEX IF NOT EXISTS idx_cost_task ON cost_records(task_id);
CREATE INDEX IF NOT EXISTS idx_cost_project ON cost_records(project_id);
CREATE INDEX IF NOT EXISTS idx_docs_project ON documents(project_id);
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON doc_chunks(document_id);
CREATE INDEX IF NOT EXISTS idx_alerts_read ON alerts(read);
"""

# ---------------------------------------------------------------------------
# 迁移：给旧表补列（V1）
# ---------------------------------------------------------------------------
def _migrate():
    conn = get_conn()
    cols = {r['name'] for r in conn.execute('PRAGMA table_info(tasks)')}
    if 'depends_on' not in cols:
        conn.execute("ALTER TABLE tasks ADD COLUMN depends_on TEXT DEFAULT '[]'")
        conn.execute('CREATE INDEX IF NOT EXISTS idx_tasks_depends ON tasks(depends_on)')
    if 'deleted' not in cols:
        conn.execute('ALTER TABLE tasks ADD COLUMN deleted INTEGER DEFAULT 0')
        conn.execute('ALTER TABLE tasks ADD COLUMN deleted_at INTEGER')
        conn.execute('CREATE INDEX IF NOT EXISTS idx_tasks_deleted ON tasks(deleted)')
    conn.commit()
    conn.close()


# ---------------------------------------------------------------------------
# V2 迁移：旧库升级（幂等）
# ---------------------------------------------------------------------------
def migrate_v2():
    """老数据库升级：V2 表由 SCHEMA 中的 CREATE TABLE IF NOT EXISTS 保证存在；
    此处处理老表缺列 / 默认数据（管理员账号、内置模板）。"""
    conn = get_conn()
    # users 表首次出现时注入默认管理员
    c = conn.execute('SELECT COUNT(*) c FROM users').fetchone()['c']
    if c == 0:
        conn.execute(
            "INSERT INTO users (username, password_hash, display_name, role, source, status, created_at) "
            "VALUES ('admin', ?, '管理员', 'admin', 'local', 'active', ?)",
            (_hash_password('admin123'), int(time.time())))
    conn.commit()
    conn.close()


def _hash_password(pwd):
    import hashlib
    return 'sha256$' + hashlib.sha256(pwd.encode('utf-8')).hexdigest()


def verify_password(pwd, pwd_hash):
    if not pwd_hash:
        return False
    if pwd_hash.startswith('sha256$'):
        import hashlib
        return hashlib.sha256(pwd.encode('utf-8')).hexdigest() == pwd_hash.split('$', 1)[1]
    return pwd == pwd_hash  # 兼容明文


def hash_password(pwd):
    return _hash_password(pwd)


def get_conn():
    conn = sqlite3.connect(DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute('PRAGMA journal_mode=WAL')
    conn.execute('PRAGMA foreign_keys=ON')
    return conn


def init_db():
    conn = get_conn()
    conn.executescript(SCHEMA)
    conn.commit()
    conn.close()
    _migrate()
    migrate_v2()


def recover_stale_runs():
    """启动恢复：进程重启后，把遗留的 running 状态标记为 failed（线程已随进程消亡）。
    覆盖：V1 任务 / V2 协作运行 / V2 评估运行。"""
    now_ts = now()
    n1 = w('UPDATE tasks SET status="failed", error="服务重启，执行中断", finished_at=? '
           'WHERE status="running"', (now_ts,))
    n2 = w('UPDATE agent_runs SET status="failed", error="服务重启，协作中断", finished_at=? '
           'WHERE status="running"', (now_ts,))
    n3 = w('UPDATE eval_runs SET status="failed", finished_at=? WHERE status="running"', (now_ts,))
    if n1 or n2 or n3:
        import logging
        logging.warning(f'recover_stale_runs: tasks={n1 or 0} agent_runs={n2 or 0} eval_runs={n3 or 0}')
    return (n1 or 0, n2 or 0, n3 or 0)


def q(sql, args=(), one=False):
    """查询"""
    conn = get_conn()
    try:
        cur = conn.execute(sql, args)
        rows = [dict(r) for r in cur.fetchall()]
        return (rows[0] if rows else None) if one else rows
    finally:
        conn.close()


def w(sql, args=()):
    """写入，返回 lastrowid"""
    conn = get_conn()
    try:
        cur = conn.execute(sql, args)
        conn.commit()
        return cur.lastrowid
    finally:
        conn.close()


def now():
    return int(time.time())


# ---------------------------------------------------------------------------
# 成本统计辅助
# ---------------------------------------------------------------------------
def monthly_worker_cost(worker_id, month_ts=None):
    """某 Worker 当月累计成本（元）"""
    if month_ts is None:
        import datetime
        month_ts = int(datetime.datetime.now().replace(day=1, hour=0, minute=0, second=0, microsecond=0).timestamp())
    rows = q(
        'SELECT COALESCE(SUM(cost),0) AS total FROM cost_records '
        'WHERE worker_id=? AND created_at>=?', (worker_id, month_ts))
    return rows[0]['total'] if rows else 0.0


def task_worker_cost(worker_id):
    """某 Worker 最近一次任务成本（用于单任务上限判断前先看历史，不作为硬限制）"""
    rows = q(
        'SELECT COALESCE(SUM(cost),0) AS total FROM cost_records WHERE worker_id=?',
        (worker_id,))
    return rows[0]['total'] if rows else 0.0


def serialize_task(t):
    t = dict(t)
    t['review_required'] = bool(t['review_required'])
    try:
        t['depends_on'] = json.loads(t.get('depends_on') or '[]')
    except Exception:
        t['depends_on'] = []
    return t


def get_setting(key, default=''):
    r = q('SELECT value FROM settings WHERE key=?', (key,), one=True)
    return r['value'] if r else default


def set_setting(key, value):
    conn = get_conn()
    try:
        conn.execute('INSERT INTO settings (key, value) VALUES (?,?) '
                     'ON CONFLICT(key) DO UPDATE SET value=excluded.value', (key, str(value)))
        conn.commit()
    finally:
        conn.close()