cmshoppe/app/db.py

"""SQLite persistence for accounts, import batches, and tasks."""

from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
import uuid
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from typing import Iterable, Optional

from . import appconfig


DEFAULT_BUSY_TIMEOUT_MS = 5000
VALID_BATCH_FIELDS = {"source_files_json", "status", "note"}
VALID_ACCOUNT_FIELDS = {
    "account_name",
    "region_host",
    "slug",
    "user_data_dir",
    "debug_port",
    "password",
    "note",
    "last_login_at",
}
PHASE_ATTEMPT_FIELDS = {
    "collect": "collect_attempts",
    "collected": "collect_attempts",
    "generate": "generate_attempts",
    "generated": "generate_attempts",
    "apply": "apply_attempts",
    "applied": "apply_attempts",
    "update": "apply_attempts",
    "updated": "apply_attempts",
}


class DbError(RuntimeError):
    """Raised when SQLite persistence cannot complete an operation."""


@dataclass(frozen=True)
class Batch:
    id: str
    source_files_json: str
    status: str
    note: Optional[str]
    created_at: str
    updated_at: str

    @property
    def source_files(self) -> list[str]:
        return json.loads(self.source_files_json)


@dataclass(frozen=True)
class Account:
    id: int
    account_name: str
    alias: str
    region_host: str
    slug: str
    user_data_dir: str
    debug_port: int
    password: Optional[str]
    note: Optional[str]
    created_at: str
    updated_at: str
    last_login_at: Optional[str]


@dataclass(frozen=True)
class Task:
    id: int
    batch_id: str
    source_file: str
    source_file_abs: str
    source_sheet: str
    source_row: int
    row_key: str
    account_name: Optional[str]
    alias: str
    item_id: str
    old_title: Optional[str]
    old_cover_path: Optional[str]
    new_title: Optional[str]
    new_cover_path: Optional[str]
    committed: int
    stage: str
    status: str
    last_error: Optional[str]
    collect_attempts: int
    generate_attempts: int
    apply_attempts: int
    imported_at: str
    collected_at: Optional[str]
    generated_at: Optional[str]
    applied_at: Optional[str]
    updated_at: str


SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS batches (
  id                TEXT PRIMARY KEY,
  source_files_json TEXT NOT NULL,
  status            TEXT NOT NULL DEFAULT 'active',
  note              TEXT,
  created_at        TEXT NOT NULL,
  updated_at        TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS accounts (
  id            INTEGER PRIMARY KEY,
  account_name  TEXT NOT NULL,
  alias         TEXT UNIQUE NOT NULL,
  region_host   TEXT NOT NULL,
  slug          TEXT UNIQUE NOT NULL,
  user_data_dir TEXT NOT NULL,
  debug_port    INTEGER NOT NULL,
  password      TEXT,
  note          TEXT,
  created_at    TEXT NOT NULL,
  updated_at    TEXT NOT NULL,
  last_login_at TEXT
);

CREATE TABLE IF NOT EXISTS tasks (
  id              INTEGER PRIMARY KEY,
  batch_id        TEXT NOT NULL REFERENCES batches(id),
  source_file     TEXT NOT NULL,
  source_file_abs TEXT NOT NULL,
  source_sheet    TEXT NOT NULL,
  source_row      INTEGER NOT NULL,
  row_key         TEXT NOT NULL UNIQUE,
  account_name    TEXT,
  alias           TEXT NOT NULL,
  item_id         TEXT NOT NULL,
  old_title       TEXT,
  old_cover_path  TEXT,
  new_title       TEXT,
  new_cover_path  TEXT,
  committed       INTEGER NOT NULL DEFAULT 0,
  stage           TEXT NOT NULL DEFAULT 'imported',
  status          TEXT NOT NULL DEFAULT 'pending',
  last_error      TEXT,
  collect_attempts  INTEGER NOT NULL DEFAULT 0,
  generate_attempts INTEGER NOT NULL DEFAULT 0,
  apply_attempts    INTEGER NOT NULL DEFAULT 0,
  imported_at     TEXT NOT NULL,
  collected_at    TEXT,
  generated_at    TEXT,
  applied_at      TEXT,
  updated_at      TEXT NOT NULL,
  UNIQUE(batch_id, source_file_abs, source_sheet, source_row)
);

CREATE INDEX IF NOT EXISTS idx_tasks_batch_stage_status
  ON tasks(batch_id, stage, status);
CREATE INDEX IF NOT EXISTS idx_tasks_alias ON tasks(alias);
CREATE INDEX IF NOT EXISTS idx_tasks_item ON tasks(item_id);
"""


def _now() -> str:
    return datetime.now().isoformat(timespec="seconds")


def _db_path(path=None) -> str:
    return path or appconfig.db_path()


def connect(path=None) -> sqlite3.Connection:
    """Open a SQLite connection with the project concurrency pragmas."""

    db_path = _db_path(path)
    if db_path != ":memory:":
        directory = os.path.dirname(os.path.abspath(db_path))
        if directory:
            os.makedirs(directory, exist_ok=True)
    conn = sqlite3.connect(db_path, timeout=DEFAULT_BUSY_TIMEOUT_MS / 1000)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}")
    conn.execute("PRAGMA synchronous=NORMAL")
    return conn


@contextmanager
def _connection(conn=None, path=None):
    owned = conn is None
    database = connect(path) if owned else conn
    try:
        yield database
    finally:
        if owned:
            database.close()


def _as_dataclass(row, cls):
    return None if row is None else cls(**dict(row))


def _fetch_one(conn, sql, params, cls):
    return _as_dataclass(conn.execute(sql, params).fetchone(), cls)


def _fetch_all(conn, sql, params, cls):
    return [_as_dataclass(row, cls) for row in conn.execute(sql, params).fetchall()]


def _validate_fields(fields, allowed):
    unknown = sorted(set(fields) - allowed)
    if unknown:
        raise DbError(f"不支持更新字段: {', '.join(unknown)}")


def _slug_from_alias(alias: str) -> str:
    base = re.sub(r"[^a-z0-9_]+", "_", alias.lower()).strip("_")
    suffix = hashlib.sha1(alias.encode("utf-8")).hexdigest()[:8]
    return f"{base}_{suffix}" if base else f"account_{suffix}"


def _row_key(batch_id, source_file_abs, source_sheet, source_row) -> str:
    return f"{batch_id}:{source_file_abs}:{source_sheet}:{source_row}"


def _attempt_field(phase: str) -> str:
    key = str(phase).lower()
    if key not in PHASE_ATTEMPT_FIELDS:
        raise DbError(f"未知阶段: {phase}")
    return PHASE_ATTEMPT_FIELDS[key]


def init_db(path=None, conn=None) -> None:
    """Create tables and indexes if they do not already exist."""

    with _connection(conn, path) as database:
        with database:
            database.executescript(SCHEMA_SQL)


def create_batch(file_paths: Iterable[str], note=None, path=None, conn=None) -> str:
    batch_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
    files = [os.path.abspath(file_path) for file_path in file_paths]
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                INSERT INTO batches
                  (id, source_files_json, status, note, created_at, updated_at)
                VALUES (?, ?, 'active', ?, ?, ?)
                """,
                (batch_id, json.dumps(files, ensure_ascii=False), note, now, now),
            )
    return batch_id


def get_batch(batch_id, path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_one(
            database,
            "SELECT * FROM batches WHERE id = ?",
            (batch_id,),
            Batch,
        )


def list_batches(status=None, path=None, conn=None):
    sql = "SELECT * FROM batches"
    params = []
    if status is not None:
        sql += " WHERE status = ?"
        params.append(status)
    sql += " ORDER BY created_at DESC, id DESC"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, Batch)


def update_batch(batch_id, path=None, conn=None, **fields) -> None:
    _validate_fields(fields, VALID_BATCH_FIELDS)
    if not fields:
        return
    fields["updated_at"] = _now()
    assignments = ", ".join(f"{field} = ?" for field in fields)
    params = list(fields.values()) + [batch_id]
    with _connection(conn, path) as database:
        with database:
            database.execute(
                f"UPDATE batches SET {assignments} WHERE id = ?",
                params,
            )


def list_accounts(path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_all(
            database,
            "SELECT * FROM accounts ORDER BY alias",
            (),
            Account,
        )


def get_account_by_alias(alias, path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_one(
            database,
            "SELECT * FROM accounts WHERE alias = ?",
            (alias,),
            Account,
        )


def add_account(
    account_name,
    alias,
    region_host,
    debug_port,
    password=None,
    note=None,
    slug=None,
    user_data_dir=None,
    path=None,
    conn=None,
):
    slug = slug or _slug_from_alias(alias)
    user_data_dir = user_data_dir or os.path.join("chrome_user_data_dir", slug)
    now = _now()
    with _connection(conn, path) as database:
        try:
            with database:
                database.execute(
                    """
                    INSERT INTO accounts
                      (account_name, alias, region_host, slug, user_data_dir,
                       debug_port, password, note, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        account_name,
                        alias,
                        region_host,
                        slug,
                        user_data_dir,
                        int(debug_port),
                        password,
                        note,
                        now,
                        now,
                    ),
                )
        except sqlite3.IntegrityError as exc:
            raise DbError(f"账号别名或 slug 已存在: {alias}") from exc
        return get_account_by_alias(alias, conn=database)


def update_account(alias, path=None, conn=None, **fields) -> None:
    _validate_fields(fields, VALID_ACCOUNT_FIELDS)
    if not fields:
        return
    fields["updated_at"] = _now()
    assignments = ", ".join(f"{field} = ?" for field in fields)
    params = list(fields.values()) + [alias]
    with _connection(conn, path) as database:
        with database:
            database.execute(
                f"UPDATE accounts SET {assignments} WHERE alias = ?",
                params,
            )


def delete_account(alias, path=None, conn=None) -> None:
    with _connection(conn, path) as database:
        with database:
            database.execute("DELETE FROM accounts WHERE alias = ?", (alias,))


def insert_tasks(batch_id, rows, path=None, conn=None) -> int:
    now = _now()
    values = []
    for row in rows:
        source_file_abs = os.path.abspath(row["source_file_abs"])
        source_sheet = row["source_sheet"]
        source_row = int(row["source_row"])
        values.append(
            (
                batch_id,
                row.get("source_file") or source_file_abs,
                source_file_abs,
                source_sheet,
                source_row,
                row.get("row_key") or _row_key(batch_id, source_file_abs, source_sheet, source_row),
                row.get("account_name"),
                row["alias"],
                str(row["item_id"]),
                now,
                now,
            )
        )
    with _connection(conn, path) as database:
        try:
            with database:
                database.executemany(
                    """
                    INSERT INTO tasks
                      (batch_id, source_file, source_file_abs, source_sheet,
                       source_row, row_key, account_name, alias, item_id,
                       imported_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    values,
                )
        except sqlite3.IntegrityError as exc:
            raise DbError("任务行重复或批次不存在") from exc
    return len(values)


def list_tasks(batch_id=None, stage=None, status=None, alias=None, path=None, conn=None):
    clauses = []
    params = []
    filters = {
        "batch_id": batch_id,
        "stage": stage,
        "status": status,
        "alias": alias,
    }
    for field, value in filters.items():
        if value is not None:
            clauses.append(f"{field} = ?")
            params.append(value)
    sql = "SELECT * FROM tasks"
    if clauses:
        sql += " WHERE " + " AND ".join(clauses)
    sql += " ORDER BY id"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, Task)


def mark_running(task_id, phase, path=None, conn=None) -> None:
    _attempt_field(phase)
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET status = 'running', last_error = NULL, updated_at = ?
                WHERE id = ?
                """,
                (_now(), int(task_id)),
            )


def mark_failed(task_id, phase, error, path=None, conn=None) -> None:
    attempt_field = _attempt_field(phase)
    with _connection(conn, path) as database:
        with database:
            database.execute(
                f"""
                UPDATE tasks
                SET status = 'failed',
                    last_error = ?,
                    {attempt_field} = {attempt_field} + 1,
                    updated_at = ?
                WHERE id = ?
                """,
                (str(error), _now(), int(task_id)),
            )


def mark_skipped(task_id, reason, path=None, conn=None) -> None:
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET status = 'skipped', last_error = ?, updated_at = ?
                WHERE id = ?
                """,
                (str(reason), _now(), int(task_id)),
            )


def set_collected(task_id, old_title, old_cover_path, path=None, conn=None) -> None:
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET old_title = ?,
                    old_cover_path = ?,
                    stage = 'collected',
                    status = 'success',
                    last_error = NULL,
                    collect_attempts = collect_attempts + 1,
                    collected_at = ?,
                    updated_at = ?
                WHERE id = ?
                """,
                (old_title, old_cover_path, now, now, int(task_id)),
            )


def set_generated(task_id, new_title, new_cover_path, path=None, conn=None) -> None:
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET new_title = ?,
                    new_cover_path = ?,
                    stage = 'generated',
                    status = 'success',
                    last_error = NULL,
                    generate_attempts = generate_attempts + 1,
                    generated_at = ?,
                    updated_at = ?
                WHERE id = ?
                """,
                (new_title, new_cover_path, now, now, int(task_id)),
            )


def set_applied(task_id, committed, error=None, path=None, conn=None) -> None:
    now = _now()
    success = bool(committed) and error is None
    with _connection(conn, path) as database:
        with database:
            if success:
                database.execute(
                    """
                    UPDATE tasks
                    SET committed = 1,
                        stage = 'applied',
                        status = 'success',
                        last_error = NULL,
                        apply_attempts = apply_attempts + 1,
                        applied_at = ?,
                        updated_at = ?
                    WHERE id = ?
                    """,
                    (now, now, int(task_id)),
                )
            else:
                database.execute(
                    """
                    UPDATE tasks
                    SET committed = 0,
                        status = 'failed',
                        last_error = ?,
                        apply_attempts = apply_attempts + 1,
                        updated_at = ?
                    WHERE id = ?
                    """,
                    (str(error or "未提交更新"), now, int(task_id)),
                )