#!/usr/bin/env python3
"""SQLite 任务持久化"""
import json
import sqlite3
import time
import uuid

from config import DB_PATH


def _conn():
    c = sqlite3.connect(DB_PATH, timeout=30)
    c.row_factory = sqlite3.Row
    return c


def init_db():
    c = _conn()
    c.execute('''
        CREATE TABLE IF NOT EXISTS tasks (
            id TEXT PRIMARY KEY,
            url TEXT NOT NULL,
            goal TEXT NOT NULL,
            status TEXT DEFAULT 'queued',
            result TEXT DEFAULT 'pending',
            max_steps INTEGER DEFAULT 30,
            timeout INTEGER DEFAULT 600,
            created_at REAL,
            started_at REAL,
            finished_at REAL,
            steps INTEGER DEFAULT 0,
            summary TEXT,
            error TEXT,
            report_path TEXT
        )
    ''')
    # 大模型配置表
    c.execute('''
        CREATE TABLE IF NOT EXISTS llm_configs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            base_url TEXT NOT NULL,
            api_key TEXT NOT NULL,
            model TEXT NOT NULL,
            vision INTEGER DEFAULT 0,
            temperature REAL DEFAULT 0.2,
            timeout INTEGER DEFAULT 120,
            is_default INTEGER DEFAULT 0,
            created_at REAL
        )
    ''')
    # tasks 表迁移：新增模型相关列（老库升级）
    cols = [r[1] for r in c.execute('PRAGMA table_info(tasks)')]
    if 'llm_config_id' not in cols:
        c.execute('ALTER TABLE tasks ADD COLUMN llm_config_id INTEGER')
    if 'llm_name' not in cols:
        c.execute('ALTER TABLE tasks ADD COLUMN llm_name TEXT')
    if 'vision' not in cols:
        c.execute('ALTER TABLE tasks ADD COLUMN vision INTEGER DEFAULT 0')
    # 首次启动：写入默认模型配置
    n = c.execute('SELECT COUNT(*) FROM llm_configs').fetchone()[0]
    if n == 0:
        from config import (LLM_BASE_URL, LLM_API_KEY, LLM_MODEL,
                            LLM_TEMPERATURE, LLM_TIMEOUT, DEFAULT_LLM_NAME)
        c.execute(
            'INSERT INTO llm_configs (name, base_url, api_key, model, vision, '
            'temperature, timeout, is_default, created_at) VALUES (?,?,?,?,?,?,?,?,?)',
            (DEFAULT_LLM_NAME, LLM_BASE_URL, LLM_API_KEY, LLM_MODEL, 0,
             LLM_TEMPERATURE, LLM_TIMEOUT, 1, time.time()))
    c.commit()
    c.close()


def create_task(url, goal, max_steps, timeout, llm_config_id=None, llm_name=None, vision=0):
    tid = uuid.uuid4().hex[:12]
    c = _conn()
    c.execute(
        'INSERT INTO tasks (id, url, goal, status, max_steps, timeout, created_at, '
        'llm_config_id, llm_name, vision) VALUES (?,?,?,?,?,?,?,?,?,?)',
        (tid, url, goal, 'queued', max_steps, timeout, time.time(),
         llm_config_id, llm_name, 1 if vision else 0))
    c.commit()
    c.close()
    return tid


def update_task(tid, **fields):
    allowed = {'status', 'result', 'started_at', 'finished_at', 'steps',
               'summary', 'error', 'report_path'}
    sets = [f'{k}=?' for k in fields if k in allowed]
    vals = [fields[k] for k in fields if k in allowed]
    if not sets:
        return
    c = _conn()
    c.execute(f'UPDATE tasks SET {", ".join(sets)} WHERE id=?', (*vals, tid))
    c.commit()
    c.close()


def get_task(tid):
    c = _conn()
    row = c.execute('SELECT * FROM tasks WHERE id=?', (tid,)).fetchone()
    c.close()
    return dict(row) if row else None


def list_tasks(limit=50):
    c = _conn()
    rows = c.execute(
        'SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?', (limit,)
    ).fetchall()
    c.close()
    return [dict(r) for r in rows]


def save_step_log(tid, step):
    """把单个步骤 JSON 追加到任务目录的 steps.jsonl"""
    from config import TASKS_DIR
    import os
    p = os.path.join(TASKS_DIR, tid, 'steps.jsonl')
    with open(p, 'a', encoding='utf-8') as f:
        f.write(json.dumps(step, ensure_ascii=False) + '\n')


def load_step_logs(tid):
    from config import TASKS_DIR
    import os
    p = os.path.join(TASKS_DIR, tid, 'steps.jsonl')
    if not os.path.exists(p):
        return []
    out = []
    with open(p, encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    out.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    return out


# ========== 大模型配置 ==========


def list_llm_configs():
    c = _conn()
    rows = c.execute(
        'SELECT * FROM llm_configs ORDER BY is_default DESC, id ASC'
    ).fetchall()
    c.close()
    return [dict(r) for r in rows]


def get_llm_config(cid):
    c = _conn()
    row = c.execute('SELECT * FROM llm_configs WHERE id=?', (cid,)).fetchone()
    c.close()
    return dict(row) if row else None


def get_default_llm_config():
    c = _conn()
    row = c.execute('SELECT * FROM llm_configs WHERE is_default=1').fetchone()
    if not row:
        row = c.execute('SELECT * FROM llm_configs ORDER BY id ASC').fetchone()
    c.close()
    return dict(row) if row else None


def create_llm_config(name, base_url, api_key, model, vision=0,
                      temperature=0.2, timeout=120, is_default=False):
    c = _conn()
    if is_default:
        c.execute('UPDATE llm_configs SET is_default=0')
    c.execute(
        'INSERT INTO llm_configs (name, base_url, api_key, model, vision, '
        'temperature, timeout, is_default, created_at) VALUES (?,?,?,?,?,?,?,?,?)',
        (name, base_url, api_key, model, 1 if vision else 0,
         temperature, timeout, 1 if is_default else 0, time.time()))
    c.commit()
    cid = c.execute('SELECT last_insert_rowid()').fetchone()[0]
    c.close()
    return cid


def update_llm_config(cid, **fields):
    allowed = {'name', 'base_url', 'api_key', 'model', 'vision',
               'temperature', 'timeout', 'is_default'}
    c = _conn()
    if fields.get('is_default'):
        c.execute('UPDATE llm_configs SET is_default=0')
    sets = []
    vals = []
    for k, v in fields.items():
        if k not in allowed:
            continue
        if k == 'vision':
            v = 1 if v else 0
        if k == 'is_default':
            v = 1 if v else 0
        sets.append(f'{k}=?')
        vals.append(v)
    if sets:
        c.execute(f'UPDATE llm_configs SET {", ".join(sets)} WHERE id=?', (*vals, cid))
    c.commit()
    c.close()


def delete_llm_config(cid):
    c = _conn()
    c.execute('DELETE FROM llm_configs WHERE id=?', (cid,))
    c.commit()
    c.close()


def set_default_llm_config(cid):
    c = _conn()
    c.execute('UPDATE llm_configs SET is_default=0')
    c.execute('UPDATE llm_configs SET is_default=1 WHERE id=?', (cid,))
    c.commit()
    c.close()


def llm_config_to_dict(row):
    """把配置行转成 llm.py 可用的 dict（不含数据库元信息）"""
    return {
        'base_url': row['base_url'],
        'api_key': row['api_key'],
        'model': row['model'],
        'temperature': row['temperature'],
        'timeout': row['timeout'],
        'vision': bool(row['vision']),
    }