# -*- 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);
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)')
    conn.commit()
    conn.close()


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


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