"""Data service for the AI image studio.""" from __future__ import annotations 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" 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"} SELECTION_TYPES = {"main", "detail"} @dataclass(frozen=True) class ImageStudioProject: id: int account_alias: str account_slug: str account_name: Optional[str] item_id: str target_main_count: int target_detail_count: int draft_prompt: 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] attempts: int points_cost: Optional[int] points_balance: Optional[int] call_id: Optional[str] 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 class ImageStudioError(RuntimeError): """Raised when the AI image studio service cannot complete an operation.""" 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 _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, target_main_count, target_detail_count, draft_prompt, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( alias, slug, name, item, 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 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 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 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(_get(project, "item_id"), "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 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 sync_original_asset_urls(project_id, image_urls, path=None, conn=None): """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"] } 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 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", path=None, conn=None, ): now = _now() task_key = str(task_key or _task_key(project_id)) 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, 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 ""), 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 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', call_id = ?, points_cost = ?, points_balance = ?, submitted_at = ?, updated_at = ? WHERE id = ? """, (str(task_id), 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, increment_attempts=False, path=None, conn=None, ): if str(status) not in JOB_STATUSES: raise db.DbError("AI工场任务状态无效") 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 = ?, 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, 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): clauses = ["status IN (?, ?)", "task_id IS NOT NULL"] params = ["submitted", "running"] 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, ): """Create/open an AI studio project and read Shopee main image URLs without editing.""" cfg = appconfig.load_config() if config is None else config 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) readiness = _ensure_account_ready_for_read( account, path=database_path, config=cfg, login_timeout=login_timeout, on_step=on_step, ) cdp = None try: _notify_step(on_step, "open_product", "start", f"商品 {item}") cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False) _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 not images: raise ImageStudioError("未读取到蝦皮商品主图 URL") assets = sync_original_asset_urls(project.id, images, path=database_path) _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)