ai-worker-platform/db.py

# -*- 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=不限
    deliver_user_id INTEGER,               -- 送达者用户 id（users.id），实时取该用户邮箱
    deliver_email TEXT DEFAULT '',         -- 送达者邮箱快照（兼容/展示，发送时以用户最新邮箱为准）
    deliver_type TEXT DEFAULT 'web',       -- 交付物类型 web=网页 / file=文件包
    deliver_note TEXT DEFAULT '',          -- 交付说明
    workspace_dir TEXT DEFAULT '',         -- 项目工作目录（相对 data/ 的目录名）
    demo_url TEXT DEFAULT '',              -- 网页交付物 Demo 访问地址
    delivered_at INTEGER,                  -- 最近一次交付/送达时间
    manager_worker_id INTEGER,             -- V3.3 AI 主管 Worker id（负责拆解/派活/监控）
    auto_status TEXT DEFAULT 'none',       -- V3.3 自动开工状态 none/running/done/failed
    auto_message TEXT DEFAULT '',          -- V3.3 自动开工最新动态
    auto_started_at INTEGER,               -- V3.3 最近一次自动开工时间
    auto_finished_at INTEGER,              -- V3.3 最近一次自动收尾时间
    review_token TEXT DEFAULT '',          -- V3.4 项目验收令牌（发给负责人，公开链接免登录验收）
    review_required INTEGER DEFAULT 1,     -- V3.4 是否需负责人验收后才算完成（默认需要）
    created_at INTEGER,
    updated_at INTEGER
);

-- V3.3 项目干活团队：AI 主管支配的多个 Worker
CREATE TABLE IF NOT EXISTS project_team_workers (
    project_id INTEGER NOT NULL,
    worker_id INTEGER NOT NULL,
    created_at INTEGER,
    PRIMARY KEY (project_id, worker_id)
);

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
);

-- V3.3 AI 主管项目级动态（拆解/派活/监控/诊断重试）
CREATE TABLE IF NOT EXISTS project_logs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    project_id INTEGER NOT NULL,
    level TEXT DEFAULT 'info',
    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 '',
    email 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 ''
);

-- ===================================================================
-- V3 表结构：交付体系（工作目录/交付物/Demo/邮件送达） + 用户授权（项目/Worker 权限）
-- ===================================================================
CREATE TABLE IF NOT EXISTS project_deliverables (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    project_id INTEGER NOT NULL,
    name TEXT NOT NULL,
    kind TEXT DEFAULT 'file',              -- file/dir/webpage/package
    path TEXT DEFAULT '',                  -- 相对项目工作目录路径 / 打包文件名
    demo_url TEXT DEFAULT '',              -- 网页交付物的 Demo 访问地址
    size INTEGER DEFAULT 0,
    note TEXT DEFAULT '',
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS user_projects (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    project_id INTEGER NOT NULL,
    perm TEXT DEFAULT 'view',              -- view 查看 / manage 管理 / admin 管理员
    created_at INTEGER,
    UNIQUE(user_id, project_id)
);

CREATE TABLE IF NOT EXISTS user_workers (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    worker_id INTEGER NOT NULL,
    perm TEXT DEFAULT 'view',              -- view 查看 / use 使用（可指派任务）/ manage 管理（可改配置）
    created_at INTEGER,
    UNIQUE(user_id, worker_id)
);

-- ===================================================================
-- V3.1 表结构：自定义角色（权限功能点）+ Worker 权限组（批量授权）
-- ===================================================================
CREATE TABLE IF NOT EXISTS roles (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL UNIQUE,
    description TEXT DEFAULT '',
    perms TEXT DEFAULT '[]',               -- JSON: 权限点 id 列表
    is_builtin INTEGER DEFAULT 0,          -- 内置角色（admin/auditor/member）不可删除
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS user_roles (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    role_id INTEGER NOT NULL,
    created_at INTEGER,
    UNIQUE(user_id, role_id)
);

CREATE TABLE IF NOT EXISTS worker_perm_groups (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS worker_perm_group_members (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    group_id INTEGER NOT NULL,
    worker_id INTEGER NOT NULL,
    perm TEXT DEFAULT 'view',              -- 组内该 Worker 权限（可覆盖组默认）
    created_at INTEGER,
    UNIQUE(group_id, worker_id)
);

CREATE TABLE IF NOT EXISTS user_worker_groups (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    group_id INTEGER NOT NULL,
    created_at INTEGER,
    UNIQUE(user_id, group_id)
);

CREATE INDEX IF NOT EXISTS idx_deliverables_project ON project_deliverables(project_id);
CREATE INDEX IF NOT EXISTS idx_user_projects_user ON user_projects(user_id);
CREATE INDEX IF NOT EXISTS idx_user_projects_project ON user_projects(project_id);
CREATE INDEX IF NOT EXISTS idx_user_workers_user ON user_workers(user_id);
CREATE INDEX IF NOT EXISTS idx_user_workers_worker ON user_workers(worker_id);
CREATE INDEX IF NOT EXISTS idx_user_roles_user ON user_roles(user_id);
CREATE INDEX IF NOT EXISTS idx_wpg_members_group ON worker_perm_group_members(group_id);
CREATE INDEX IF NOT EXISTS idx_uwg_user ON user_worker_groups(user_id);
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)')
    # V3：projects 交付字段
    pcols = {r['name'] for r in conn.execute('PRAGMA table_info(projects)')}
    for col, ddl in (
        ('deliver_email', "ALTER TABLE projects ADD COLUMN deliver_email TEXT DEFAULT ''"),
        ('deliver_type', "ALTER TABLE projects ADD COLUMN deliver_type TEXT DEFAULT 'web'"),
        ('deliver_note', "ALTER TABLE projects ADD COLUMN deliver_note TEXT DEFAULT ''"),
        ('workspace_dir', "ALTER TABLE projects ADD COLUMN workspace_dir TEXT DEFAULT ''"),
        ('demo_url', "ALTER TABLE projects ADD COLUMN demo_url TEXT DEFAULT ''"),
        ('delivered_at', 'ALTER TABLE projects ADD COLUMN delivered_at INTEGER'),
    ):
        if col not in pcols:
            conn.execute(ddl)
    # V3.2：projects 送达者改为用户 id（实时取邮箱）
    if 'deliver_user_id' not in pcols:
        conn.execute('ALTER TABLE projects ADD COLUMN deliver_user_id INTEGER')
    # V3.3：projects AI 主管 + 自动开工状态
    for col, ddl in (
        ('manager_worker_id', 'ALTER TABLE projects ADD COLUMN manager_worker_id INTEGER'),
        ('auto_status', "ALTER TABLE projects ADD COLUMN auto_status TEXT DEFAULT 'none'"),
        ('auto_message', "ALTER TABLE projects ADD COLUMN auto_message TEXT DEFAULT ''"),
        ('auto_started_at', 'ALTER TABLE projects ADD COLUMN auto_started_at INTEGER'),
        ('auto_finished_at', 'ALTER TABLE projects ADD COLUMN auto_finished_at INTEGER'),
    ):
        if col not in pcols:
            conn.execute(ddl)
    # V3.4：负责人验收
    pcols = {r['name'] for r in conn.execute('PRAGMA table_info(projects)')}
    for col, ddl in (
        ('review_token', "ALTER TABLE projects ADD COLUMN review_token TEXT DEFAULT ''"),
        ('review_required', 'ALTER TABLE projects ADD COLUMN review_required INTEGER DEFAULT 1'),
    ):
        if col not in pcols:
            conn.execute(ddl)
    conn.execute('''CREATE TABLE IF NOT EXISTS project_team_workers (
        project_id INTEGER NOT NULL,
        worker_id INTEGER NOT NULL,
        created_at INTEGER,
        PRIMARY KEY (project_id, worker_id)
    )''')
    conn.execute('''CREATE TABLE IF NOT EXISTS project_logs (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        project_id INTEGER NOT NULL,
        level TEXT DEFAULT 'info',
        message TEXT DEFAULT '',
        created_at INTEGER
    )''')
    # V3.2：users 邮箱列 + 存量用户默认邮箱 + 旧项目按邮箱回填送达者用户
    ucols = {r['name'] for r in conn.execute('PRAGMA table_info(users)')}
    if 'email' not in ucols:
        conn.execute("ALTER TABLE users ADD COLUMN email TEXT DEFAULT ''")
        for r in conn.execute('SELECT id, username, email FROM users WHERE email IS NULL OR email=\'\''):
            conn.execute('UPDATE users SET email=? WHERE id=?',
                         (f'{r["username"]}@tphai.com', r['id']))
    # 旧项目：按 deliver_email 匹配用户回填 deliver_user_id
    for r in conn.execute("SELECT id, deliver_email FROM projects WHERE (deliver_user_id IS NULL OR deliver_user_id=0) "
                          "AND deliver_email IS NOT NULL AND deliver_email != ''"):
        u = conn.execute('SELECT id FROM users WHERE email=?', (r['deliver_email'],)).fetchone()
        if u:
            conn.execute('UPDATE projects SET deliver_user_id=? WHERE id=?', (u['id'], r['id']))
    # V3：老项目补齐工作目录名
    for r in conn.execute("SELECT id, workspace_dir FROM projects WHERE workspace_dir IS NULL OR workspace_dir=''"):
        conn.execute('UPDATE projects SET workspace_dir=? WHERE id=?',
                     ('project_%d' % r['id'], r['id']))
    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()
    seed_builtin_roles()


# V3.1 内置角色权限点定义（admin 为特殊值 ALL，表示全部权限）
BUILTIN_ROLE_PERMS = {
    'admin': ['*'],
    # 审计员：全量只读
    'auditor': ['dashboard.view', 'project.view', 'worker.view', 'agent.view', 'eval.view',
                'template.view', 'report.view', 'log.view', 'alert.view'],
    # 成员：基础功能（资源级仍受项目/Worker 授权约束）
    'member': ['dashboard.view', 'project.view', 'project.create', 'project.deliver',
               'worker.view', 'agent.view', 'agent.run', 'eval.view', 'eval.run',
               'template.view', 'report.view', 'log.view', 'alert.view'],
}


def seed_builtin_roles():
    """幂等：仅首次创建 admin/auditor/member 内置角色（不覆盖管理员后续编辑）"""
    import json as _json
    conn = get_conn()
    try:
        for name, perms in BUILTIN_ROLE_PERMS.items():
            r = conn.execute('SELECT id FROM roles WHERE name=?', (name,)).fetchone()
            if not r:
                conn.execute(
                    'INSERT INTO roles (name, description, perms, is_builtin, created_at, updated_at) '
                    'VALUES (?,?,?,1,?,?)',
                    (name, {'admin': '超级管理员：全部权限', 'auditor': '审计员：全量只读',
                            'member': '成员：基础功能，资源级按项目/Worker 授权'}[name],
                     _json.dumps(perms), now(), now()))
        conn.commit()
    finally:
        conn.close()


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 serialize_task_light(t):
    """轻量序列化（列表接口用）：去掉大字段 output_text / error，
    看板、DAG、回收站等列表视图不需要它们，可显著减小传输体积。
    详情接口 /api/tasks/<id> 仍返回完整数据。"""
    t = serialize_task(t)
    t.pop('output_text', None)
    t.pop('error', None)
    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()