cmshoppe/app/image_studio.py

"""Data service for the AI image studio."""

from __future__ import annotations

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

from . import accounts, appconfig, chrome, db, editor
from .config import make_slug


PROJECT_STATUS_ACTIVE = "active"
PROJECT_BINDING_DRAFT = "draft"
PROJECT_BINDING_BOUND = "bound"
PROJECT_BINDING_STATES = {PROJECT_BINDING_DRAFT, PROJECT_BINDING_BOUND}
TEMPORARY_ITEM_PREFIX = "draft_"
ASSET_KIND_ORIGINAL = "original"
ASSET_STATUS_AVAILABLE = "available"
ASSET_STATUS_MISSING = "missing"
ASSET_STATUSES = {ASSET_STATUS_AVAILABLE, ASSET_STATUS_MISSING}
JOB_STATUSES = {"pending", "submitted", "running", "succeeded", "failed", "expired", "cancelled"}
JOB_RESUMABLE_STATUSES = {"submitted", "running"}
JOB_RECOVERY_NONE = "none"
JOB_RECOVERY_RESUME = "resume"
JOB_RECOVERY_REGENERATE = "regenerate"
JOB_RECOVERY_ACTIONS = {
    JOB_RECOVERY_NONE,
    JOB_RECOVERY_RESUME,
    JOB_RECOVERY_REGENERATE,
}
SELECTION_TYPES = {"main", "detail"}


@dataclass(frozen=True)
class ImageStudioProject:
    id: int
    account_alias: str
    account_slug: str
    account_name: Optional[str]
    item_id: str
    storage_key: str
    binding_state: str
    target_main_count: int
    target_detail_count: int
    draft_prompt: Optional[str]
    suite_settings_json: str
    current_generation_round_key: Optional[str]
    status: str
    created_at: str
    updated_at: str
    deleted_at: Optional[str]
    deleted_reason: Optional[str]


@dataclass(frozen=True)
class ImageStudioAsset:
    id: int
    project_id: int
    kind: str
    remote_url: Optional[str]
    local_path: Optional[str]
    aspect_ratio: Optional[str]
    parent_asset_id: Optional[int]
    prompt: Optional[str]
    status: str
    source_order: int
    created_at: str
    updated_at: str


@dataclass(frozen=True)
class ImageStudioJob:
    id: int
    project_id: int
    source_asset_id: Optional[int]
    output_asset_id: Optional[int]
    generation_source: str
    provider: str
    job_type: str
    task_key: str
    task_id: Optional[str]
    status: str
    prompt: Optional[str]
    error: Optional[str]
    recovery_action: str
    attempts: int
    points_cost: Optional[int]
    points_balance: Optional[int]
    call_id: Optional[str]
    generation_round_key: Optional[str]
    generation_slot_index: Optional[int]
    created_at: str
    updated_at: str
    submitted_at: Optional[str]
    finished_at: Optional[str]


@dataclass(frozen=True)
class ImageStudioSelection:
    id: int
    project_id: int
    selection_type: str
    position: int
    asset_id: int
    created_at: str
    updated_at: str


@dataclass(frozen=True)
class ImageStudioGenerationRound:
    project_id: int
    generation_round_key: Optional[str]
    created_at: Optional[str]
    updated_at: Optional[str]
    job_count: int
    slot_count: int
    succeeded_count: int
    failed_count: int
    cancelled_count: int
    active_count: int
    retry_count: int
    is_current: bool
    is_legacy: bool


class ImageStudioError(RuntimeError):
    """Raised when the AI image studio service cannot complete an operation."""


class ImageStudioProjectConflictError(ImageStudioError):
    """Raised when a draft cannot be bound because the formal project already exists."""


class ImageStudioPullCancelled(ImageStudioError):
    """Raised when a read-only Shopee image pull stops at a safe boundary."""

    def __init__(self, project=None, assets=None):
        super().__init__("用户停止拉取蝦皮主图")
        self.project = project
        self.assets = list(assets or [])


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


@contextmanager
def _connection(conn=None, path=None):
    with db._connection(conn, path) as database:
        yield database


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


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


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


def _get(obj, name, default=None):
    if obj is None:
        return default
    if isinstance(obj, dict):
        return obj.get(name, default)
    return getattr(obj, name, default)


def _safe_component(value, default):
    text = str(value or "").strip()
    if not text:
        text = str(default)
    safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in text).strip("_")
    return safe or str(default)


def _account_fields(account=None, account_alias=None, account_name=None, account_slug=None):
    alias = str(account_alias or _get(account, "alias") or "").strip()
    if not alias:
        raise db.DbError("AI工场项目缺少账号别名")
    name = str(account_name or _get(account, "account_name") or "").strip() or None
    slug = str(account_slug or _get(account, "slug") or "").strip() or make_slug(alias)
    return alias, name, _safe_component(slug, "account")


def _task_key(project_id):
    return f"image-studio-{int(project_id)}-{uuid.uuid4().hex}"


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


def _normalize_item_id(item_id) -> str:
    text = str(item_id or "").strip()
    if not text:
        raise ImageStudioError("商品ID不能为空")
    return text


def is_formal_item_id(item_id) -> bool:
    return str(item_id or "").strip().isdigit()


def is_draft_project(project) -> bool:
    return str(_get(project, "binding_state", "")).strip() == PROJECT_BINDING_DRAFT


def project_storage_key(project) -> str:
    return str(_get(project, "storage_key") or _get(project, "item_id") or "").strip()


def _new_draft_item_id() -> str:
    return TEMPORARY_ITEM_PREFIX + uuid.uuid4().hex


def _notify_step(callback, step, result="start", detail=None):
    if callback is None:
        return
    payload = {"step": step, "result": result}
    if detail:
        payload["detail"] = str(detail)
    try:
        callback(payload)
    except Exception:
        pass


def _is_definitive_logged_out(status):
    reason = str((status or {}).get("reason") or "").strip()
    url = str((status or {}).get("url") or "").lower()
    return reason.startswith("LOGIN_PAGE") or (
        "accounts.shopee." in url and "/seller/login" in url
    )


def _login_status_detail(status):
    status = status or {}
    reason = status.get("reason") or "未知原因"
    url = status.get("url") or "未知URL"
    cookie_names = [str(name) for name in (status.get("cookie_names") or []) if name]
    cookie_text = ",".join(sorted(cookie_names)) if cookie_names else "未读到登录Cookie"
    return f"原因={reason}，URL={url}，Cookie名称={cookie_text}"


def _ensure_account_ready_for_read(account, *, path=None, config=None, login_timeout=8, on_step=None):
    port = accounts.normalize_debug_port(_get(account, "debug_port"))
    alias = str(_get(account, "alias") or "").strip()
    _notify_step(on_step, "ensure_chrome", "start", f"账号 {alias} debug_port={port}")
    if chrome.is_running(port):
        chrome_info = {
            "action": "reused",
            "reused": True,
            "launched": False,
            "debug_port": port,
        }
        _notify_step(on_step, "ensure_chrome", "reused", f"账号 {alias} Chrome 已打开")
    else:
        try:
            chrome_info = accounts.launch_for_login(account, path=path, config=config)
        except Exception as exc:
            raise ImageStudioError(f"账号 {alias} Chrome 启动失败：{exc}") from exc
        action = "launched" if chrome_info.get("launched") else "reused"
        _notify_step(on_step, "ensure_chrome", action, f"账号 {alias} Chrome 已{ '启动' if action == 'launched' else '复用' }")

    _notify_step(on_step, "login_check", "start", f"账号 {alias}")
    try:
        status = accounts.detect_login(account, timeout=login_timeout, path=path, config=config)
    except Exception as exc:
        status = {"logged_in": False, "reason": f"LOGIN_CHECK_FAILED: {exc}"}
    if status.get("logged_in"):
        _notify_step(on_step, "login_check", "success", f"账号 {alias} 已登录")
        return {**chrome_info, "login_status": status, "login_uncertain": False}
    if _is_definitive_logged_out(status):
        detail = _login_status_detail(status)
        _notify_step(on_step, "login_check", "blocked", detail)
        raise ImageStudioError(f"账号 {alias} 未登录：{detail}")
    detail = _login_status_detail(status)
    _notify_step(on_step, "login_check", "uncertain", f"账号 {alias} 登录状态暂不稳定，继续尝试读取商品：{detail}")
    return {**chrome_info, "login_status": status, "login_uncertain": True}


def _ensure_assets_belong_to_project(database, project_id, asset_ids):
    ids = [int(asset_id) for asset_id in asset_ids if asset_id is not None]
    if not ids:
        return
    placeholders = ",".join("?" for _ in ids)
    rows = database.execute(
        f"SELECT id FROM image_studio_assets WHERE project_id = ? AND id IN ({placeholders})",
        [int(project_id), *ids],
    ).fetchall()
    found = {int(row["id"]) for row in rows}
    missing = sorted(set(ids) - found)
    if missing:
        raise db.DbError("AI工场资产不属于当前项目")


def get_project(project_id, path=None, conn=None, include_deleted=False):
    clauses = ["id = ?"]
    params = [int(project_id)]
    if not include_deleted:
        clauses.append("deleted_at IS NULL")
    sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses)
    with _connection(conn, path) as database:
        return _fetch_one(database, sql, params, ImageStudioProject)


def get_project_by_account_item(account_alias, item_id, path=None, conn=None, include_deleted=False):
    clauses = ["account_alias = ?", "item_id = ?"]
    params = [str(account_alias).strip(), str(item_id).strip()]
    if not include_deleted:
        clauses.append("deleted_at IS NULL")
    sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses)
    with _connection(conn, path) as database:
        return _fetch_one(database, sql, params, ImageStudioProject)


def create_or_get_project(
    account=None,
    item_id=None,
    *,
    account_alias=None,
    account_name=None,
    account_slug=None,
    target_main_count=9,
    target_detail_count=12,
    draft_prompt="",
    path=None,
    conn=None,
):
    alias, name, slug = _account_fields(
        account,
        account_alias=account_alias,
        account_name=account_name,
        account_slug=account_slug,
    )
    item = str(item_id or _get(account, "item_id") or "").strip()
    if not item:
        raise db.DbError("AI工场项目缺少商品ID")
    now = _now()
    with _connection(conn, path) as database:
        existing = get_project_by_account_item(alias, item, conn=database, include_deleted=True)
        if existing is not None:
            if existing.deleted_at is not None:
                with database:
                    database.execute(
                        """
                        UPDATE image_studio_projects
                        SET deleted_at = NULL,
                            deleted_reason = NULL,
                            status = ?,
                            updated_at = ?
                        WHERE id = ?
                        """,
                        (PROJECT_STATUS_ACTIVE, now, existing.id),
                    )
                return get_project(existing.id, conn=database)
            return existing
        with database:
            cursor = database.execute(
                """
                INSERT INTO image_studio_projects
                  (account_alias, account_slug, account_name, item_id, storage_key, binding_state,
                   target_main_count, target_detail_count, draft_prompt,
                   status, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    alias,
                    slug,
                    name,
                    item,
                    item,
                    PROJECT_BINDING_BOUND,
                    int(target_main_count),
                    int(target_detail_count),
                    str(draft_prompt or ""),
                    PROJECT_STATUS_ACTIVE,
                    now,
                    now,
                ),
            )
            project_id = cursor.lastrowid
        return get_project(project_id, conn=database)


def create_draft_project(
    account=None,
    *,
    account_alias=None,
    account_name=None,
    account_slug=None,
    target_main_count=9,
    target_detail_count=12,
    draft_prompt="",
    path=None,
    conn=None,
):
    """Create one project-bound local draft without exposing its internal item key."""

    alias, name, slug = _account_fields(
        account,
        account_alias=account_alias,
        account_name=account_name,
        account_slug=account_slug,
    )
    item = _new_draft_item_id()
    now = _now()
    with _connection(conn, path) as database:
        with database:
            cursor = database.execute(
                """
                INSERT INTO image_studio_projects
                  (account_alias, account_slug, account_name, item_id, storage_key, binding_state,
                   target_main_count, target_detail_count, draft_prompt,
                   status, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    alias,
                    slug,
                    name,
                    item,
                    item,
                    PROJECT_BINDING_DRAFT,
                    int(target_main_count),
                    int(target_detail_count),
                    str(draft_prompt or ""),
                    PROJECT_STATUS_ACTIVE,
                    now,
                    now,
                ),
            )
            project_id = cursor.lastrowid
        return get_project(project_id, conn=database)


def bind_draft_project(project_id, item_id, path=None, conn=None):
    """Bind one active draft to a formal numeric item ID without moving its files."""

    item = str(item_id or "").strip()
    if not is_formal_item_id(item):
        raise ImageStudioError("正式商品ID必须是数字")
    with _connection(conn, path) as database:
        source = get_project(project_id, conn=database, include_deleted=True)
        if source is None:
            raise ImageStudioError("临时草稿不存在")
        if source.deleted_at is not None:
            raise ImageStudioError("临时草稿已删除，无法绑定商品")
        if not is_draft_project(source):
            if source.item_id == item:
                return source
            raise ImageStudioError("当前项目不是临时草稿，不能重新绑定商品")
        existing = get_project_by_account_item(
            source.account_alias,
            item,
            conn=database,
            include_deleted=True,
        )
        if existing is not None and existing.id != source.id:
            raise ImageStudioProjectConflictError("该商品项目已存在，不能覆盖或合并")
        now = _now()
        with database:
            database.execute(
                """
                UPDATE image_studio_projects
                SET item_id = ?, binding_state = ?, updated_at = ?
                WHERE id = ? AND deleted_at IS NULL
                """,
                (item, PROJECT_BINDING_BOUND, now, int(source.id)),
            )
        return get_project(source.id, conn=database)


def list_projects(path=None, conn=None, include_deleted=False):
    sql = "SELECT * FROM image_studio_projects"
    if not include_deleted:
        sql += " WHERE deleted_at IS NULL"
    sql += " ORDER BY updated_at DESC, id DESC"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, (), ImageStudioProject)


def list_recoverable_draft_projects(path=None, conn=None):
    """Return non-empty drafts that should reappear as product-suite task tabs."""

    sql = """
        SELECT p.*
        FROM image_studio_projects AS p
        WHERE p.deleted_at IS NULL
          AND p.binding_state = ?
          AND (
            TRIM(COALESCE(p.draft_prompt, '')) <> ''
            OR
            EXISTS (
              SELECT 1 FROM image_studio_assets AS a
              WHERE a.project_id = p.id
            )
            OR EXISTS (
              SELECT 1 FROM image_studio_jobs AS j
              WHERE j.project_id = p.id
            )
          )
        ORDER BY p.updated_at DESC, p.id DESC
    """
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, (PROJECT_BINDING_DRAFT,), ImageStudioProject)


def project_has_content(project_id, path=None, conn=None) -> bool:
    with _connection(conn, path) as database:
        row = database.execute(
            """
            SELECT EXISTS(
              SELECT 1 FROM image_studio_assets WHERE project_id = ?
            ) OR EXISTS(
              SELECT 1 FROM image_studio_jobs WHERE project_id = ?
            ) OR EXISTS(
              SELECT 1 FROM image_studio_projects
              WHERE id = ? AND TRIM(COALESCE(draft_prompt, '')) <> ''
            ) AS has_content
            """,
            (int(project_id), int(project_id), int(project_id)),
        ).fetchone()
        return bool(row["has_content"] if row is not None else False)


def discard_empty_draft_project(project_id, reason="空临时草稿", path=None, conn=None):
    """Soft-delete a newly created, still empty draft and leave user files untouched."""

    with _connection(conn, path) as database:
        project = get_project(project_id, conn=database, include_deleted=True)
        if project is None or project.deleted_at is not None or not is_draft_project(project):
            return project
        if project_has_content(project.id, conn=database):
            return project
        return soft_delete_project(project.id, reason=reason, conn=database)


def update_project_prompt(project_id, draft_prompt, path=None, conn=None):
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE image_studio_projects
                SET draft_prompt = ?, updated_at = ?
                WHERE id = ? AND deleted_at IS NULL
                """,
                (str(draft_prompt or ""), _now(), int(project_id)),
            )
        return get_project(project_id, conn=database)


def project_suite_settings(project) -> dict:
    raw = _get(project, "suite_settings_json", "{}") or "{}"
    try:
        value = json.loads(str(raw))
    except (TypeError, ValueError):
        return {}
    return value if isinstance(value, dict) else {}


def update_project_suite_settings(project_id, settings, path=None, conn=None):
    if not isinstance(settings, dict):
        raise db.DbError("商品套图设置必须是对象")
    payload = json.dumps(settings, ensure_ascii=False, sort_keys=True)
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE image_studio_projects
                SET suite_settings_json = ?, updated_at = ?
                WHERE id = ? AND deleted_at IS NULL
                """,
                (payload, _now(), int(project_id)),
            )
        return get_project(project_id, conn=database)


def soft_delete_project(project_id, reason="", path=None, conn=None):
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE image_studio_projects
                SET deleted_at = ?, deleted_reason = ?, updated_at = ?
                WHERE id = ? AND deleted_at IS NULL
                """,
                (now, str(reason or ""), now, int(project_id)),
            )
        return get_project(project_id, conn=database, include_deleted=True)


def project_image_dirs(image_root, project):
    root = os.path.abspath(
        os.path.join(
            str(image_root or "images"),
            "pool",
            _safe_component(_get(project, "account_slug"), "account"),
            _safe_component(project_storage_key(project), "item"),
        )
    )
    return {
        "root": root,
        "originals": os.path.join(root, "originals"),
        "generated": os.path.join(root, "generated"),
        "exports": os.path.join(root, "exports"),
    }


def default_project_image_dirs(project, config=None):
    return project_image_dirs(appconfig.image_dir(config), project)


def add_asset(
    project_id,
    kind,
    *,
    remote_url=None,
    local_path=None,
    aspect_ratio=None,
    parent_asset_id=None,
    prompt=None,
    status=ASSET_STATUS_AVAILABLE,
    source_order=0,
    path=None,
    conn=None,
):
    now = _now()
    abs_local_path = os.path.abspath(local_path) if local_path else None
    with _connection(conn, path) as database:
        with database:
            cursor = database.execute(
                """
                INSERT INTO image_studio_assets
                  (project_id, kind, remote_url, local_path, aspect_ratio,
                   parent_asset_id, prompt, status, source_order, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    int(project_id),
                    str(kind),
                    remote_url,
                    abs_local_path,
                    aspect_ratio,
                    parent_asset_id,
                    prompt,
                    str(status or ASSET_STATUS_AVAILABLE),
                    int(source_order or 0),
                    now,
                    now,
                ),
            )
            asset_id = cursor.lastrowid
        return get_asset(asset_id, conn=database)


def get_asset(asset_id, path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_one(
            database,
            "SELECT * FROM image_studio_assets WHERE id = ?",
            (int(asset_id),),
            ImageStudioAsset,
        )


def list_assets(project_id, kind=None, include_missing=True, path=None, conn=None):
    clauses = ["project_id = ?"]
    params = [int(project_id)]
    if kind is not None:
        clauses.append("kind = ?")
        params.append(str(kind))
    if not include_missing:
        clauses.append("status != ?")
        params.append(ASSET_STATUS_MISSING)
    sql = "SELECT * FROM image_studio_assets WHERE " + " AND ".join(clauses)
    sql += " ORDER BY source_order, id"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, ImageStudioAsset)


def reorder_original_assets(project_id, asset_ids, path=None, conn=None):
    ordered_ids = [int(asset_id) for asset_id in asset_ids]
    if len(ordered_ids) != len(set(ordered_ids)):
        raise db.DbError("商品原图排序不能包含重复图片")
    with _connection(conn, path) as database:
        rows = database.execute(
            """
            SELECT id FROM image_studio_assets
            WHERE project_id = ? AND kind = ?
            ORDER BY source_order, id
            """,
            (int(project_id), ASSET_KIND_ORIGINAL),
        ).fetchall()
        existing_ids = [int(row["id"]) for row in rows]
        if set(existing_ids) != set(ordered_ids):
            raise db.DbError("商品原图排序必须包含当前项目的全部原图")
        now = _now()
        with database:
            for source_order, asset_id in enumerate(ordered_ids, 1):
                database.execute(
                    """
                    UPDATE image_studio_assets
                    SET source_order = ?, updated_at = ?
                    WHERE id = ? AND project_id = ? AND kind = ?
                    """,
                    (
                        source_order,
                        now,
                        asset_id,
                        int(project_id),
                        ASSET_KIND_ORIGINAL,
                    ),
                )
        return list_assets(
            project_id,
            kind=ASSET_KIND_ORIGINAL,
            path=path,
            conn=database,
        )


def asset_reference_counts(asset_id, path=None, conn=None):
    """Return selection/job references for one asset before pool removal."""

    asset_id = int(asset_id)
    with _connection(conn, path) as database:
        selection_count = database.execute(
            "SELECT COUNT(*) FROM image_studio_selections WHERE asset_id = ?",
            (asset_id,),
        ).fetchone()[0]
        source_job_count = database.execute(
            "SELECT COUNT(*) FROM image_studio_jobs WHERE source_asset_id = ?",
            (asset_id,),
        ).fetchone()[0]
        output_job_count = database.execute(
            "SELECT COUNT(*) FROM image_studio_jobs WHERE output_asset_id = ?",
            (asset_id,),
        ).fetchone()[0]
    return {
        "selection": int(selection_count),
        "source_job": int(source_job_count),
        "output_job": int(output_job_count),
        "total": int(selection_count + source_job_count + output_job_count),
    }


def remove_asset_if_unused(asset_id, path=None, conn=None):
    """Remove one pool asset row only when no job or selection references it.

    The local image file is intentionally kept on disk. AI工场 removal is a
    pool-level operation, not a destructive file cleanup.
    """

    asset_id = int(asset_id)
    with _connection(conn, path) as database:
        asset = get_asset(asset_id, conn=database)
        if asset is None:
            raise db.DbError("照片不存在")
        counts = asset_reference_counts(asset_id, conn=database)
        if counts["total"]:
            raise db.DbError("照片正在被生成任务或终选引用，不能移除")
        with database:
            database.execute("DELETE FROM image_studio_assets WHERE id = ?", (asset_id,))
        return asset


def remove_original_assets_if_unused(project_id, asset_ids, path=None, conn=None):
    """Atomically remove unreferenced original assets from one project.

    Local files are intentionally retained. If any requested asset is invalid or
    referenced, no asset row is removed.
    """

    project_id = int(project_id)
    ordered_ids = []
    seen = set()
    for value in asset_ids or []:
        asset_id = int(value)
        if asset_id not in seen:
            seen.add(asset_id)
            ordered_ids.append(asset_id)
    if not ordered_ids:
        raise db.DbError("请选择要删除的商品原图")

    placeholders = ",".join("?" for _ in ordered_ids)
    with _connection(conn, path) as database:
        with database:
            project = get_project(project_id, conn=database)
            if project is None:
                raise db.DbError("商品套图项目不存在或已删除")
            rows = database.execute(
                f"""
                SELECT * FROM image_studio_assets
                WHERE project_id = ? AND kind = ? AND id IN ({placeholders})
                """,
                [project_id, ASSET_KIND_ORIGINAL, *ordered_ids],
            ).fetchall()
            by_id = {int(row["id"]): row for row in rows}
            if set(by_id) != set(ordered_ids):
                raise db.DbError("选中的商品原图不存在或不属于当前项目")

            referenced_ids = [
                asset_id
                for asset_id in ordered_ids
                if asset_reference_counts(asset_id, conn=database)["total"]
            ]
            if referenced_ids:
                raise db.DbError("选中的图片正在被生成任务或终选引用，不能移除")

            database.execute(
                f"""
                DELETE FROM image_studio_assets
                WHERE project_id = ? AND kind = ? AND id IN ({placeholders})
                """,
                [project_id, ASSET_KIND_ORIGINAL, *ordered_ids],
            )
            remaining_rows = database.execute(
                """
                SELECT id FROM image_studio_assets
                WHERE project_id = ? AND kind = ?
                ORDER BY source_order, id
                """,
                (project_id, ASSET_KIND_ORIGINAL),
            ).fetchall()
            now = _now()
            for source_order, row in enumerate(remaining_rows, 1):
                database.execute(
                    """
                    UPDATE image_studio_assets
                    SET source_order = ?, updated_at = ?
                    WHERE id = ? AND project_id = ? AND kind = ?
                    """,
                    (
                        source_order,
                        now,
                        int(row["id"]),
                        project_id,
                        ASSET_KIND_ORIGINAL,
                    ),
                )

        return [_row_to_dataclass(by_id[asset_id], ImageStudioAsset) for asset_id in ordered_ids]


def restore_original_asset_snapshot(project_id, snapshots, path=None, conn=None):
    """Restore status/order for original assets that existed before a pull."""

    project_id = int(project_id)
    normalized = []
    seen = set()
    for value in snapshots or []:
        asset_id = int(_get(value, "id"))
        if asset_id in seen:
            continue
        status = str(_get(value, "status") or ASSET_STATUS_AVAILABLE)
        if status not in ASSET_STATUSES:
            raise db.DbError("商品原图快照状态无效")
        normalized.append(
            {
                "id": asset_id,
                "status": status,
                "source_order": max(0, int(_get(value, "source_order") or 0)),
            }
        )
        seen.add(asset_id)
    if not normalized:
        return []
    ids = [item["id"] for item in normalized]
    placeholders = ",".join("?" for _ in ids)
    with _connection(conn, path) as database:
        with database:
            rows = database.execute(
                f"""
                SELECT id FROM image_studio_assets
                WHERE project_id = ? AND kind = ? AND id IN ({placeholders})
                """,
                [project_id, ASSET_KIND_ORIGINAL, *ids],
            ).fetchall()
            if {int(row["id"]) for row in rows} != set(ids):
                raise db.DbError("拉取前商品原图快照已失效")
            now = _now()
            for item in normalized:
                database.execute(
                    """
                    UPDATE image_studio_assets
                    SET status = ?, source_order = ?, updated_at = ?
                    WHERE id = ? AND project_id = ? AND kind = ?
                    """,
                    (
                        item["status"],
                        item["source_order"],
                        now,
                        item["id"],
                        project_id,
                        ASSET_KIND_ORIGINAL,
                    ),
                )
            remaining_rows = database.execute(
                f"""
                SELECT id FROM image_studio_assets
                WHERE project_id = ? AND kind = ?
                  AND id NOT IN ({placeholders})
                ORDER BY source_order, id
                """,
                [project_id, ASSET_KIND_ORIGINAL, *ids],
            ).fetchall()
            next_order = max(
                [int(item["source_order"]) for item in normalized] + [0]
            ) + 1
            for row in remaining_rows:
                database.execute(
                    """
                    UPDATE image_studio_assets
                    SET source_order = ?, updated_at = ?
                    WHERE id = ? AND project_id = ? AND kind = ?
                    """,
                    (
                        next_order,
                        now,
                        int(row["id"]),
                        project_id,
                        ASSET_KIND_ORIGINAL,
                    ),
                )
                next_order += 1
        return [
            get_asset(item["id"], conn=database)
            for item in normalized
        ]


def sync_original_asset_urls(project_id, image_urls, path=None, conn=None, max_assets=16):
    """Store the read-only Shopee main image URL snapshot as remote-only assets."""

    normalized = []
    seen_urls = set()
    for position, item in enumerate(image_urls or [], start=1):
        if isinstance(item, dict):
            src = str(item.get("src") or "").strip()
            source_order = int(item.get("index") or position)
        else:
            src = str(item or "").strip()
            source_order = position
        if not src or src in seen_urls:
            continue
        seen_urls.add(src)
        normalized.append({"src": src, "source_order": source_order})

    now = _now()
    with _connection(conn, path) as database:
        with database:
            existing_rows = database.execute(
                """
                SELECT * FROM image_studio_assets
                WHERE project_id = ? AND kind = ?
                """,
                (int(project_id), ASSET_KIND_ORIGINAL),
            ).fetchall()
            by_url = {
                str(row["remote_url"]): row
                for row in existing_rows
                if row["remote_url"]
            }
            local_only_count = sum(
                1
                for row in existing_rows
                if not row["remote_url"] and row["status"] != ASSET_STATUS_MISSING
            )
            remote_limit = max(0, int(max_assets) - local_only_count)
            normalized = normalized[:remote_limit]
            active_ids = set()
            for item in normalized:
                row = by_url.get(item["src"])
                if row is not None:
                    active_ids.add(int(row["id"]))
                    database.execute(
                        """
                        UPDATE image_studio_assets
                        SET source_order = ?,
                            status = ?,
                            updated_at = ?
                        WHERE id = ?
                        """,
                        (
                            int(item["source_order"]),
                            ASSET_STATUS_AVAILABLE,
                            now,
                            int(row["id"]),
                        ),
                    )
                else:
                    cursor = database.execute(
                        """
                        INSERT INTO image_studio_assets
                          (project_id, kind, remote_url, local_path, status,
                           source_order, created_at, updated_at)
                        VALUES (?, ?, ?, NULL, ?, ?, ?, ?)
                        """,
                        (
                            int(project_id),
                            ASSET_KIND_ORIGINAL,
                            item["src"],
                            ASSET_STATUS_AVAILABLE,
                            int(item["source_order"]),
                            now,
                            now,
                        ),
                    )
                    active_ids.add(int(cursor.lastrowid))
            missing_ids = [
                int(row["id"])
                for row in existing_rows
                if row["remote_url"] and int(row["id"]) not in active_ids
            ]
            if missing_ids:
                placeholders = ",".join("?" for _ in missing_ids)
                database.execute(
                    f"""
                    UPDATE image_studio_assets
                    SET status = ?, updated_at = ?
                    WHERE id IN ({placeholders})
                    """,
                    [ASSET_STATUS_MISSING, now, *missing_ids],
                )
        return list_assets(project_id, kind=ASSET_KIND_ORIGINAL, conn=database)


def mark_asset_status(asset_id, status, path=None, conn=None):
    if str(status) not in ASSET_STATUSES:
        raise db.DbError("AI工场资产状态必须是 available 或 missing")
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE image_studio_assets
                SET status = ?, updated_at = ?
                WHERE id = ?
                """,
                (str(status), _now(), int(asset_id)),
            )
        return get_asset(asset_id, conn=database)


def update_asset_local_path(asset_id, local_path, status=ASSET_STATUS_AVAILABLE, path=None, conn=None):
    if str(status) not in ASSET_STATUSES:
        raise db.DbError("AI工场资产状态必须是 available 或 missing")
    abs_local_path = os.path.abspath(local_path) if local_path else None
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE image_studio_assets
                SET local_path = ?, status = ?, updated_at = ?
                WHERE id = ?
                """,
                (abs_local_path, str(status), _now(), int(asset_id)),
            )
        return get_asset(asset_id, conn=database)


def create_job(
    project_id,
    *,
    source_asset_id=None,
    job_type="main",
    prompt="",
    task_key=None,
    generation_source="cmhub",
    provider="cmhub",
    generation_round_key=None,
    generation_slot_index=None,
    path=None,
    conn=None,
):
    now = _now()
    task_key = str(task_key or _task_key(project_id))
    if generation_round_key is not None:
        generation_round_key = str(generation_round_key).strip()
        if not generation_round_key:
            raise db.DbError("AI工场生成轮次标识不能为空")
    if generation_slot_index is not None:
        try:
            generation_slot_index = int(generation_slot_index)
        except (TypeError, ValueError) as exc:
            raise db.DbError("AI工场生成槽位编号无效") from exc
        if generation_slot_index < 0:
            raise db.DbError("AI工场生成槽位编号不能小于0")
    with _connection(conn, path) as database:
        try:
            with database:
                _ensure_assets_belong_to_project(database, project_id, [source_asset_id])
                cursor = database.execute(
                    """
                    INSERT INTO image_studio_jobs
                      (project_id, source_asset_id, generation_source, provider,
                       job_type, task_key, status, prompt, generation_round_key,
                       generation_slot_index, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?)
                    """,
                    (
                        int(project_id),
                        source_asset_id,
                        str(generation_source or "cmhub"),
                        str(provider or "cmhub"),
                        str(job_type),
                        task_key,
                        str(prompt or ""),
                        generation_round_key,
                        generation_slot_index,
                        now,
                        now,
                    ),
                )
                job_id = cursor.lastrowid
        except Exception as exc:
            raise db.DbError(f"创建AI工场生图任务失败: {exc}") from exc
        return get_job(job_id, conn=database)


def get_job(job_id, path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_one(
            database,
            "SELECT * FROM image_studio_jobs WHERE id = ?",
            (int(job_id),),
            ImageStudioJob,
        )


def list_jobs(project_id, *, statuses=None, path=None, conn=None):
    clauses = ["project_id = ?"]
    params = [int(project_id)]
    requested_statuses = [str(status) for status in (statuses or []) if str(status)]
    if requested_statuses:
        invalid = set(requested_statuses) - JOB_STATUSES
        if invalid:
            raise db.DbError("图片生成任务状态无效")
        placeholders = ",".join("?" for _ in requested_statuses)
        clauses.append(f"status IN ({placeholders})")
        params.extend(requested_statuses)
    sql = "SELECT * FROM image_studio_jobs WHERE " + " AND ".join(clauses)
    sql += " ORDER BY updated_at DESC, id DESC"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, ImageStudioJob)


def get_current_generation_round(project_id, path=None, conn=None):
    project = get_project(project_id, path=path, conn=conn)
    if project is None:
        return None
    value = project.current_generation_round_key
    return str(value).strip() if value else None


def set_current_generation_round(project_id, generation_round_key, path=None, conn=None):
    project_id = int(project_id)
    round_key = None
    if generation_round_key is not None:
        round_key = str(generation_round_key).strip()
        if not round_key:
            raise db.DbError("当前生成轮次标识不能为空")
    with _connection(conn, path) as database:
        with database:
            project = get_project(project_id, conn=database)
            if project is None:
                raise db.DbError("AI工场商品项目不存在或已删除")
            if round_key is not None:
                row = database.execute(
                    """
                    SELECT 1 FROM image_studio_jobs
                    WHERE project_id = ? AND generation_round_key = ?
                    LIMIT 1
                    """,
                    (project_id, round_key),
                ).fetchone()
                if row is None:
                    raise db.DbError("当前生成轮次不属于该商品项目")
            database.execute(
                """
                UPDATE image_studio_projects
                SET current_generation_round_key = ?, updated_at = ?
                WHERE id = ? AND deleted_at IS NULL
                """,
                (round_key, _now(), project_id),
            )
        return get_project(project_id, conn=database)


def promote_generation_round_if_success(project_id, generation_round_key, path=None, conn=None):
    project_id = int(project_id)
    round_key = str(generation_round_key or "").strip()
    if not round_key:
        raise db.DbError("当前生成轮次标识不能为空")
    with _connection(conn, path) as database:
        with database:
            project = get_project(project_id, conn=database)
            if project is None:
                raise db.DbError("AI工场商品项目不存在或已删除")
            success = database.execute(
                """
                SELECT 1 FROM image_studio_jobs
                WHERE project_id = ?
                  AND generation_round_key = ?
                  AND status = 'succeeded'
                LIMIT 1
                """,
                (project_id, round_key),
            ).fetchone()
            if success is None:
                return False
            database.execute(
                """
                UPDATE image_studio_projects
                SET current_generation_round_key = ?, updated_at = ?
                WHERE id = ? AND deleted_at IS NULL
                """,
                (round_key, _now(), project_id),
            )
        return True


def list_generation_rounds(project_id, *, limit=None, offset=0, path=None, conn=None):
    project_id = int(project_id)
    try:
        offset = max(0, int(offset or 0))
    except (TypeError, ValueError) as exc:
        raise db.DbError("生成轮次偏移量无效") from exc
    if limit is not None:
        try:
            limit = int(limit)
        except (TypeError, ValueError) as exc:
            raise db.DbError("生成轮次数量无效") from exc
        if limit <= 0:
            return []
    current_round_key = get_current_generation_round(project_id, path=path, conn=conn)
    sql = """
        SELECT generation_round_key,
               MIN(created_at) AS created_at,
               MAX(updated_at) AS updated_at,
               COUNT(*) AS job_count,
               COUNT(DISTINCT generation_slot_index) AS slot_count,
               SUM(CASE WHEN status = 'succeeded' THEN 1 ELSE 0 END) AS succeeded_count,
               SUM(CASE WHEN status IN ('failed', 'expired') THEN 1 ELSE 0 END) AS failed_count,
               SUM(CASE WHEN status = 'cancelled' THEN 1 ELSE 0 END) AS cancelled_count,
               SUM(CASE WHEN status IN ('pending', 'submitted', 'running') THEN 1 ELSE 0 END) AS active_count
        FROM image_studio_jobs
        WHERE project_id = ?
        GROUP BY generation_round_key
        ORDER BY CASE WHEN generation_round_key IS NULL THEN 1 ELSE 0 END,
                 MAX(id) DESC
    """
    params = [project_id]
    if limit is not None:
        sql += " LIMIT ? OFFSET ?"
        params.extend([limit, offset])
    elif offset:
        sql += " LIMIT -1 OFFSET ?"
        params.append(offset)
    with _connection(conn, path) as database:
        rows = database.execute(sql, params).fetchall()
    rounds = []
    for row in rows:
        round_key = row["generation_round_key"]
        job_count = int(row["job_count"] or 0)
        slot_count = int(row["slot_count"] or 0)
        rounds.append(
            ImageStudioGenerationRound(
                project_id=project_id,
                generation_round_key=round_key,
                created_at=row["created_at"],
                updated_at=row["updated_at"],
                job_count=job_count,
                slot_count=slot_count,
                succeeded_count=int(row["succeeded_count"] or 0),
                failed_count=int(row["failed_count"] or 0),
                cancelled_count=int(row["cancelled_count"] or 0),
                active_count=int(row["active_count"] or 0),
                retry_count=max(0, job_count - slot_count) if round_key else 0,
                is_current=bool(round_key and round_key == current_round_key),
                is_legacy=round_key is None,
            )
        )
    return rounds


def list_generation_round_current_jobs(project_id, generation_round_key, path=None, conn=None):
    project_id = int(project_id)
    if generation_round_key is None:
        sql = """
            SELECT * FROM image_studio_jobs
            WHERE project_id = ? AND generation_round_key IS NULL
            ORDER BY created_at ASC, id ASC
        """
        params = (project_id,)
    else:
        round_key = str(generation_round_key).strip()
        if not round_key:
            raise db.DbError("生成轮次标识不能为空")
        sql = """
            SELECT jobs.*
            FROM image_studio_jobs AS jobs
            INNER JOIN (
                SELECT generation_slot_index, MAX(id) AS latest_id
                FROM image_studio_jobs
                WHERE project_id = ? AND generation_round_key = ?
                GROUP BY generation_slot_index
            ) AS latest ON latest.latest_id = jobs.id
            WHERE jobs.project_id = ? AND jobs.generation_round_key = ?
            ORDER BY jobs.generation_slot_index ASC, jobs.id ASC
        """
        params = (project_id, round_key, project_id, round_key)
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, ImageStudioJob)


def list_generation_round_attempts(project_id, generation_round_key, path=None, conn=None):
    project_id = int(project_id)
    if generation_round_key is None:
        sql = """
            SELECT * FROM image_studio_jobs
            WHERE project_id = ? AND generation_round_key IS NULL
            ORDER BY created_at ASC, id ASC
        """
        params = (project_id,)
    else:
        round_key = str(generation_round_key).strip()
        if not round_key:
            raise db.DbError("生成轮次标识不能为空")
        sql = """
            SELECT * FROM image_studio_jobs
            WHERE project_id = ? AND generation_round_key = ?
            ORDER BY generation_slot_index ASC, id ASC
        """
        params = (project_id, round_key)
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, ImageStudioJob)


def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points_balance=None, path=None, conn=None):
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE image_studio_jobs
                SET task_id = ?,
                    status = 'submitted',
                    recovery_action = ?,
                    call_id = ?,
                    points_cost = ?,
                    points_balance = ?,
                    submitted_at = ?,
                    updated_at = ?
                WHERE id = ?
                """,
                (
                    str(task_id),
                    JOB_RECOVERY_RESUME,
                    call_id,
                    points_cost,
                    points_balance,
                    now,
                    now,
                    int(job_id),
                ),
            )
        return get_job(job_id, conn=database)


def update_job_status(
    job_id,
    status,
    *,
    error=None,
    output_asset_id=None,
    points_balance=None,
    recovery_action=None,
    increment_attempts=False,
    path=None,
    conn=None,
):
    if str(status) not in JOB_STATUSES:
        raise db.DbError("AI工场任务状态无效")
    if recovery_action is not None and str(recovery_action) not in JOB_RECOVERY_ACTIONS:
        raise db.DbError("AI工场任务恢复方式无效")
    if recovery_action is None and str(status) == "succeeded":
        recovery_action = JOB_RECOVERY_NONE
    now = _now()
    terminal = str(status) in {"succeeded", "failed", "expired", "cancelled"}
    with _connection(conn, path) as database:
        with database:
            if output_asset_id is not None:
                job = get_job(job_id, conn=database)
                if job is None:
                    raise db.DbError("AI工场任务不存在")
                _ensure_assets_belong_to_project(database, job.project_id, [output_asset_id])
            database.execute(
                """
                UPDATE image_studio_jobs
                SET status = ?,
                    error = ?,
                    recovery_action = COALESCE(?, recovery_action),
                    output_asset_id = COALESCE(?, output_asset_id),
                    points_balance = COALESCE(?, points_balance),
                    attempts = attempts + ?,
                    finished_at = CASE WHEN ? THEN ? ELSE finished_at END,
                    updated_at = ?
                WHERE id = ?
                """,
                (
                    str(status),
                    error,
                    recovery_action,
                    output_asset_id,
                    points_balance,
                    1 if increment_attempts else 0,
                    1 if terminal else 0,
                    now,
                    now,
                    int(job_id),
                ),
            )
        return get_job(job_id, conn=database)


def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_downloads=False):
    if include_failed_downloads:
        clauses = [
            "task_id IS NOT NULL",
            "recovery_action = ?",
            "status IN (?, ?, ?, ?)",
        ]
        params = [JOB_RECOVERY_RESUME, "submitted", "running", "failed", "cancelled"]
    else:
        clauses = ["status IN (?, ?)", "task_id IS NOT NULL", "recovery_action = ?"]
        params = ["submitted", "running", JOB_RECOVERY_RESUME]
    if project_id is not None:
        clauses.append("project_id = ?")
        params.append(int(project_id))
    sql = "SELECT * FROM image_studio_jobs WHERE " + " AND ".join(clauses)
    sql += " ORDER BY updated_at, id"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, ImageStudioJob)


def replace_selections(project_id, selection_type, asset_ids: Iterable[int], path=None, conn=None):
    selection = str(selection_type)
    if selection not in SELECTION_TYPES:
        raise db.DbError("终选类型必须是 main 或 detail")
    ids = [int(asset_id) for asset_id in asset_ids]
    if len(ids) != len(set(ids)):
        raise db.DbError("同一资产不能重复加入同一类终选")
    now = _now()
    with _connection(conn, path) as database:
        with database:
            _ensure_assets_belong_to_project(database, project_id, ids)
            database.execute(
                """
                DELETE FROM image_studio_selections
                WHERE project_id = ? AND selection_type = ?
                """,
                (int(project_id), selection),
            )
            database.executemany(
                """
                INSERT INTO image_studio_selections
                  (project_id, selection_type, position, asset_id, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                [
                    (int(project_id), selection, index + 1, asset_id, now, now)
                    for index, asset_id in enumerate(ids)
                ],
            )
        return list_selections(project_id, selection, conn=database)


def list_selections(project_id, selection_type=None, path=None, conn=None):
    clauses = ["project_id = ?"]
    params = [int(project_id)]
    if selection_type is not None:
        clauses.append("selection_type = ?")
        params.append(str(selection_type))
    sql = "SELECT * FROM image_studio_selections WHERE " + " AND ".join(clauses)
    sql += " ORDER BY selection_type, position"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, ImageStudioSelection)


def pull_remote_main_image_urls(
    account_or_alias,
    item_id,
    *,
    path=None,
    config=None,
    login_timeout=8,
    on_step=None,
    should_stop=None,
):
    """Create/open an AI studio project and read Shopee main image URLs without editing."""

    cfg = appconfig.load_config() if config is None else config
    should_stop = should_stop or (lambda: False)
    if should_stop():
        raise ImageStudioPullCancelled()
    database_path = _db_path(path, cfg)
    db.init_db(database_path)
    item = _normalize_item_id(item_id)
    try:
        account = accounts.resolve_account(account_or_alias, path=database_path, config=cfg)
    except Exception as exc:
        raise ImageStudioError(f"AI工场账号不可用：{exc}") from exc

    project = create_or_get_project(account, item_id=item, path=database_path)
    if should_stop():
        raise ImageStudioPullCancelled(project=project)
    readiness = _ensure_account_ready_for_read(
        account,
        path=database_path,
        config=cfg,
        login_timeout=login_timeout,
        on_step=on_step,
    )
    if should_stop():
        raise ImageStudioPullCancelled(project=project)
    cdp = None
    try:
        if should_stop():
            raise ImageStudioPullCancelled(project=project)
        _notify_step(on_step, "open_product", "start", f"商品 {item}")
        cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False)
        if should_stop():
            raise ImageStudioPullCancelled(project=project)
        _notify_step(on_step, "open_product", "success", f"商品 {item}")
        _notify_step(on_step, "read_main_images", "start", f"商品 {item}")
        images = editor.read_product_image_urls(cdp)
        if should_stop():
            raise ImageStudioPullCancelled(project=project)
        if not images:
            raise ImageStudioError("未读取到蝦皮商品主图 URL")
        assets = sync_original_asset_urls(project.id, images, path=database_path)
        if should_stop():
            raise ImageStudioPullCancelled(project=project, assets=assets)
        _notify_step(on_step, "read_main_images", "success", f"读取 {len(images)} 张主图 URL")
        project = get_project(project.id, path=database_path)
        return {
            "project": project,
            "images": images,
            "assets": assets,
            "account": account,
            "readiness": readiness,
        }
    except ImageStudioError:
        raise
    except editor.EditorError as exc:
        raise ImageStudioError(str(exc)) from exc
    except Exception as exc:
        raise ImageStudioError(f"读取蝦皮原主图失败：{exc}") from exc
    finally:
        if cdp is not None:
            editor.close_readonly_product(cdp)