"""SQLite 持久化任务队列（重启不丢，借鉴 OpenRig durable queue）。

设计：
- 单文件 SQLite（WAL 模式）
- 任务状态机: pending → running → done | failed | timeout
- 幂等：同一 task_id 重复入队只保留一条
- 原子领取：BEGIN IMMEDIATE + UPDATE ... WHERE status='pending'
"""
from __future__ import annotations

import json
import sqlite3
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# 幂等窗口：契约 R-1 固化为 **24h**（同 idem_key 24h 内只执行一次）
IDEM_TTL_S = 24 * 3600

SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
    id          TEXT PRIMARY KEY,
    assignee    TEXT NOT NULL,
    prompt      TEXT NOT NULL,
    context     TEXT DEFAULT '{}',
    status      TEXT NOT NULL DEFAULT 'pending',
    result      TEXT,
    error       TEXT,
    created_at  REAL NOT NULL,
    claimed_at  REAL,
    finished_at REAL,
    deadline_s  REAL DEFAULT 600,
    idem_key    TEXT UNIQUE
);
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_tasks_assignee ON tasks(assignee, status);

CREATE TABLE IF NOT EXISTS rig_meta (
    key   TEXT PRIMARY KEY,
    value TEXT
);
"""


@dataclass
class Task:
    id: str
    assignee: str
    prompt: str
    context: dict[str, Any] = field(default_factory=dict)
    status: str = "pending"
    result: str | None = None
    error: str | None = None
    created_at: float = field(default_factory=time.time)
    claimed_at: float | None = None
    finished_at: float | None = None
    deadline_s: float = 600.0
    idem_key: str | None = None

    def to_row(self) -> tuple:
        return (
            self.id, self.assignee, self.prompt, json.dumps(self.context, ensure_ascii=False),
            self.status, self.result, self.error,
            self.created_at, self.claimed_at, self.finished_at, self.deadline_s,
        )

    @classmethod
    def from_row(cls, row: tuple) -> Task:
        return cls(
            id=row[0], assignee=row[1], prompt=row[2],
            context=json.loads(row[3] or "{}"), status=row[4],
            result=row[5], error=row[6], created_at=row[7],
            claimed_at=row[8], finished_at=row[9], deadline_s=row[10],
            idem_key=(row[11] if len(row) > 11 else None),   # G1：SCHEMA 第 12 列；老库迁移后同样存在
        )


def ensure_idem_column(conn) -> None:
    """G1：老库补列。SQLite `ADD COLUMN` 不能带 UNIQUE ⇒ 用**唯一索引**实现同等约束；
    历史行 idem_key 为 NULL（NULL 不违反唯一索引），不影响既有数据。"""
    cols = [r[1] for r in conn.execute("PRAGMA table_info(tasks)").fetchall()]
    if "idem_key" not in cols:
        conn.execute("ALTER TABLE tasks ADD COLUMN idem_key TEXT")
    conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_tasks_idem ON tasks(idem_key)")
    conn.commit()


class TaskQueue:
    """SQLite 持久化任务队列。线程安全（sqlite3 check_same_thread=False + 锁）。"""

    def __init__(self, db_path: str | Path = ":memory:"):
        self.db_path = str(db_path)
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self.conn.execute("PRAGMA journal_mode=WAL")
        self.conn.execute("PRAGMA busy_timeout=5000")
        self.conn.executescript(SCHEMA)
        ensure_idem_column(self.conn)      # G1：老库补列（新库已在 SCHEMA 里）
        self.conn.commit()

    def close(self):
        self.conn.close()

    # ── 入队 ──

    def enqueue(self, assignee: str, prompt: str,
                context: dict | None = None,
                task_id: str | None = None,
                deadline_s: float = 600.0,
                idem_key: str | None = None) -> Task:
        """入队。**G1 幂等**：给了 `idem_key` 且 24h 内已有同键任务 ⇒ 直接返回那条，**不重复执行**
        （对齐契约 R-1 / PRD A4「重发三次只执行一次」）。"""
        if idem_key:
            row = self.conn.execute(
                "SELECT id FROM tasks WHERE idem_key=? AND created_at > ?",
                (idem_key, time.time() - IDEM_TTL_S)).fetchone()
            if row:
                got = self.get(row[0])
                if got is not None:
                    return got
        tid = task_id or f"task-{uuid.uuid4().hex[:12]}"
        task = Task(id=tid, assignee=assignee, prompt=prompt,
                    context=context or {}, deadline_s=deadline_s, idem_key=idem_key)
        # 幂等：已存在则跳过
        cur = self.conn.execute("SELECT id FROM tasks WHERE id=?", (tid,))
        if cur.fetchone():
            return self.get(tid)  # type: ignore
        self.conn.execute(
            "INSERT INTO tasks (id,assignee,prompt,context,status,created_at,deadline_s,idem_key) "
            "VALUES (?,?,?,?,?,?,?,?)",
            (tid, assignee, prompt, json.dumps(context or {}, ensure_ascii=False),
             "pending", task.created_at, deadline_s, idem_key),
        )
        self.conn.commit()
        return task

    # ── 领取（原子） ──

    def claim(self, assignee: str | None = None, task_id: str | None = None) -> Task | None:
        """原子领取一个 pending 任务。assignee=None 时不限制 assignee。"""
        self.conn.execute("BEGIN IMMEDIATE")
        try:
            if task_id:
                cur = self.conn.execute(
                    "SELECT * FROM tasks WHERE id=? AND status='pending'", (task_id,))
            else:
                where = "status='pending'"
                params: list[Any] = []
                if assignee:
                    where += " AND assignee=?"
                    params.append(assignee)
                cur = self.conn.execute(
                    f"SELECT * FROM tasks WHERE {where} ORDER BY created_at LIMIT 1", params)
            row = cur.fetchone()
            if not row:
                self.conn.execute("COMMIT")
                return None
            task = Task.from_row(row)
            self.conn.execute(
                "UPDATE tasks SET status='running', claimed_at=? WHERE id=?",
                (time.time(), task.id))
            self.conn.execute("COMMIT")
            task.status = "running"
            task.claimed_at = time.time()
            return task
        except Exception:
            self.conn.execute("ROLLBACK")
            raise

    # ── 完成 ──

    def complete(self, task_id: str, result: str | None = None):
        self.conn.execute(
            "UPDATE tasks SET status='done', result=?, finished_at=? WHERE id=?",
            (result, time.time(), task_id))
        self.conn.commit()

    def fail(self, task_id: str, error: str):
        self.conn.execute(
            "UPDATE tasks SET status='failed', error=?, finished_at=? WHERE id=?",
            (error, time.time(), task_id))
        self.conn.commit()

    def timeout(self, task_id: str):
        self.conn.execute(
            "UPDATE tasks SET status='timeout', finished_at=? WHERE id=?",
            (time.time(), task_id))
        self.conn.commit()

    # ── 查询 ──

    def get(self, task_id: str) -> Task | None:
        cur = self.conn.execute("SELECT * FROM tasks WHERE id=?", (task_id,))
        row = cur.fetchone()
        return Task.from_row(row) if row else None

    def list(self, status: str | None = None, assignee: str | None = None,
             limit: int = 50) -> list[Task]:
        where, params = ["1=1"], []
        if status:
            where.append("status=?")
            params.append(status)
        if assignee:
            where.append("assignee=?")
            params.append(assignee)
        cur = self.conn.execute(
            f"SELECT * FROM tasks WHERE {' AND '.join(where)} ORDER BY created_at DESC LIMIT ?",
            params + [limit])
        return [Task.from_row(r) for r in cur.fetchall()]

    def count(self, status: str | None = None) -> int:
        if status:
            cur = self.conn.execute("SELECT COUNT(*) FROM tasks WHERE status=?", (status,))
        else:
            cur = self.conn.execute("SELECT COUNT(*) FROM tasks")
        return cur.fetchone()[0]

    def pending_for(self, assignee: str) -> int:
        cur = self.conn.execute(
            "SELECT COUNT(*) FROM tasks WHERE assignee=? AND status='pending'", (assignee,))
        return cur.fetchone()[0]

    # ── 超时扫描 ──

    def sweep_timeouts(self) -> list[str]:
        """把超时的 running 任务标为 timeout。返回被标记的 task_id 列表。"""
        now = time.time()
        cur = self.conn.execute(
            "SELECT id FROM tasks WHERE status='running' AND claimed_at IS NOT NULL "
            "AND (claimed_at + deadline_s) < ?", (now,))
        ids = [r[0] for r in cur.fetchall()]
        for tid in ids:
            self.timeout(tid)
        return ids

    # ── meta ──

    def set_meta(self, key: str, value: str):
        self.conn.execute(
            "INSERT OR REPLACE INTO rig_meta (key,value) VALUES (?,?)", (key, value))
        self.conn.commit()

    def get_meta(self, key: str) -> str | None:
        cur = self.conn.execute("SELECT value FROM rig_meta WHERE key=?", (key,))
        row = cur.fetchone()
        return row[0] if row else None