cmshoppe/app/image_studio.py

"""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 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)