cmshoppe/app/ai.py

"""AI generation helpers backed by configurable HTTP model endpoints."""

import base64
import ipaddress
from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
import copy
import json
import mimetypes
import os
import socket
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid

import requests

from . import appconfig, db, diagnostics, image_paths
from . import prompts as prompt_module



class AIError(RuntimeError):
    """Raised when AI generation cannot complete."""

class CMHubError(AIError):
    """Structured cmhub gateway error."""

    def __init__(self, code, message, status=None, retryable=False, retry_after=None):
        self.code = str(code or "unknown")
        self.status = status
        self.retryable = bool(retryable)
        self.retry_after = retry_after
        super().__init__("cmhub %s: %s" % (self.code, message))


CMHUB_IMAGE_MAX_BYTES = 20 * 1024 * 1024

_RESOLUTION_SIZES = {
    "512": (512, 512),
    "1k": (1024, 1024),
    "2k": (2048, 2048),
    "4k": (4096, 4096),
}


def gen_title(
    title_prompt,
    old_title,
    retry=None,
    config=None,
    models_path=appconfig.AI_MODELS_PATH,
    on_step=None,
    on_event=None,
    cmhub_config_path=appconfig.CMHUB_CONFIG_PATH,
):
    """Generate a new product title from a prompt and the old title."""

    cfg = appconfig.load_config() if config is None else config
    ai_cfg = appconfig.ai_config(cfg)
    if _ai_backend(ai_cfg) == "cmhub":
        return _gen_title_cmhub(
            title_prompt,
            old_title,
            retry=retry,
            config=cfg,
            cmhub_config_path=cmhub_config_path,
            on_step=on_step,
            on_event=on_event,
        )
    _notify_step(on_step, "load_text_model")
    model = _role_model("text", ai_cfg.get("default_text_model"), models_path)
    _notify_step(on_step, "title_build_request")
    payload = _chat_payload(
        model,
        [
            {"role": "system", "content": str(title_prompt or "").strip()},
            {
                "role": "user",
                "content": "旧标题：\n%s\n\n请只返回新标题，不要解释。" % str(old_title or ""),
            },
        ],
    )
    attempts = _attempt_count(ai_cfg, retry)
    _notify_step(on_step, "title_request")
    data = _call_with_retry(
        model,
        payload,
        cfg,
        attempts,
        request_kind="json",
        on_retry=lambda attempt, total_attempts, exc: _notify_retry(
            on_step,
            "title_request",
            attempt,
            total_attempts,
            exc,
            model,
        ),
    )
    _notify_step(on_step, "title_parse_response")
    text = _extract_text(data).strip()
    if not text:
        raise AIError("AI 返回为空标题")
    return text


def gen_cover(
    cover_prompt,
    old_cover_path,
    out_path,
    resolution=None,
    jpg_quality=None,
    retry=None,
    config=None,
    models_path=appconfig.AI_MODELS_PATH,
    on_step=None,
    on_event=None,
    cmhub_config_path=appconfig.CMHUB_CONFIG_PATH,
):
    """Generate a new cover image and save it as a JPEG file."""

    _notify_step(on_step, "cover_validate_input")
    old_cover_path = os.path.abspath(str(old_cover_path))
    if not os.path.exists(old_cover_path):
        raise FileNotFoundError("旧封面图片不存在: %s" % old_cover_path)
    if not out_path:
        raise AIError("缺少新封面输出路径")

    cfg = appconfig.load_config() if config is None else config
    ai_cfg = appconfig.ai_config(cfg)
    if _ai_backend(ai_cfg) == "cmhub":
        return _gen_cover_cmhub(
            cover_prompt,
            old_cover_path,
            out_path,
            resolution=resolution,
            jpg_quality=jpg_quality,
            retry=retry,
            config=cfg,
            cmhub_config_path=cmhub_config_path,
            on_step=on_step,
            on_event=on_event,
        )
    _notify_step(on_step, "load_image_model")
    model = _role_model("image", ai_cfg.get("default_image_model"), models_path)
    resolution = str(resolution or ai_cfg.get("resolution", "1k"))
    quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90))
    attempts = _attempt_count(ai_cfg, retry)

    api_type = model.get("api_type", "auto")
    _notify_step(on_step, "cover_build_request")
    if api_type == "images_edits":
        body, content_type = _image_edit_body(model, cover_prompt, old_cover_path, resolution)
        _notify_step(on_step, "cover_request")
        data = _call_with_retry(
            model,
            body,
            cfg,
            attempts,
            request_kind="multipart",
            content_type=content_type,
            on_retry=lambda attempt, total_attempts, exc: _notify_retry(
                on_step,
                "cover_request",
                attempt,
                total_attempts,
                exc,
                model,
            ),
        )
    else:
        payload = _image_chat_payload(model, cover_prompt, old_cover_path, resolution)
        _notify_step(on_step, "cover_request")
        data = _call_with_retry(
            model,
            payload,
            cfg,
            attempts,
            request_kind="json",
            on_retry=lambda attempt, total_attempts, exc: _notify_retry(
                on_step,
                "cover_request",
                attempt,
                total_attempts,
                exc,
                model,
            ),
        )

    _notify_step(on_step, "cover_parse_response")
    image_bytes = _extract_image_bytes(data, model, cfg)
    _notify_step(on_step, "cover_save")
    return _save_jpeg(image_bytes, out_path, resolution, quality)


def _notify_step(callback, step):
    if callback is None:
        return
    try:
        callback(step)
    except Exception:
        pass



def _notify_retry(callback, step, attempt, attempts, exc, model):
    if callback is None:
        return
    try:
        callback(
            {
                "step": step,
                "result": "retry",
                "attempt": attempt,
                "attempts": attempts,
                "detail": _redact(str(exc), model),
            }
        )
    except Exception:
        pass

def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=None):
    """Generate titles first, then covers, and persist each successful task."""

    runtime = dict(ai_cfg or {})
    config = _runtime_config(runtime)
    generation_cfg = appconfig.ai_config(config)
    generation_cfg.update(
        {
            key: value
            for key, value in runtime.items()
            if key in {
                "title_concurrency",
                "image_concurrency",
                "retry",
                "jpg_quality",
                "resolution",
                "generate_cover",
            }
        }
    )
    db_path = runtime.get("db_path")
    models_path = runtime.get("models_path", appconfig.AI_MODELS_PATH)
    cmhub_config_path = runtime.get("cmhub_config_path", appconfig.CMHUB_CONFIG_PATH)
    image_root = runtime.get("image_dir") or appconfig.image_dir(config)
    account_by_alias = runtime.get("account_by_alias") or {}
    on_task_update = runtime.get("on_task_update")
    on_event = runtime.get("on_event")
    on_error = runtime.get("on_error")
    title_prompt = _prompt_value(prompts, "title")
    cover_prompt = _prompt_value(prompts, "cover")
    generate_cover = bool(generation_cfg.get("generate_cover", False))
    should_stop = should_stop or (lambda: False)
    eligible = [
        task for task in list(tasks)
        if getattr(task, "stage", None) == "collected"
    ]
    summary = {
        "ok": True,
        "total": len(eligible),
        "title_done": 0,
        "cover_done": 0,
        "cover_total": len(eligible) if generate_cover else 0,
        "generated_done": 0,
        "failed": 0,
        "cancelled": False,
        "generate_cover": generate_cover,
    }
    _emit_generation_progress(on_progress, summary)
    title_results = {}
    step_by_task = {}
    step_lock = threading.Lock()

    def set_step(task, step):
        with step_lock:
            step_by_task[getattr(task, "id", None)] = str(step)

    def get_step(task, fallback):
        with step_lock:
            return step_by_task.get(getattr(task, "id", None), fallback)

    def step_callback(task, phase):
        def callback(event):
            if isinstance(event, dict):
                step = event.get("step") or "unknown"
                result = event.get("result") or "start"
                set_step(task, step)
                _emit_generation_event(
                    on_event,
                    task,
                    phase,
                    step,
                    result,
                    detail=event.get("detail"),
                    level=event.get("level") or ("warning" if result == "retry" else "info"),
                    attempt=event.get("attempt"),
                    attempts=event.get("attempts"),
                    metadata=event.get("metadata"),
                )
                return
            set_step(task, event)
            _emit_generation_event(on_event, task, phase, event, "start")
        return callback

    with ThreadPoolExecutor(
        max_workers=max(1, int(generation_cfg.get("title_concurrency", 1)))
    ) as executor:
        futures = {}
        for task in eligible:
            if should_stop():
                summary["cancelled"] = True
                break
            set_step(task, "title_submit")
            _emit_generation_event(on_event, task, "title", "title_submit", "start")
            futures[
                executor.submit(
                    gen_title,
                    title_prompt,
                    getattr(task, "old_title", "") or "",
                    retry=generation_cfg.get("retry"),
                    config=config,
                    models_path=models_path,
                    on_step=step_callback(task, "title"),
                    on_event=step_callback(task, "title"),
                    cmhub_config_path=cmhub_config_path,
                )
            ] = task
        for future in as_completed(futures):
            task = futures[future]
            if should_stop():
                summary["cancelled"] = True
                _cancel_pending(futures)
            try:
                new_title = future.result()
                title_results[task.id] = new_title
                summary["title_done"] += 1
                set_step(task, "title_done")
                _emit_generation_event(on_event, task, "title", "title_done", "success")
                if not generate_cover:
                    set_step(task, "db_write")
                    _emit_generation_event(on_event, task, "title", "db_write", "start")
                    db.set_generated(task.id, new_title, None, path=db_path)
                    summary["generated_done"] += 1
                    if on_task_update is not None:
                        on_task_update(
                            task.id,
                            {
                                "stage": "generated",
                                "status": "success",
                                "new_title": new_title,
                                "new_cover_path": None,
                            },
                        )
                    _emit_generation_event(
                        on_event,
                        task,
                        "title",
                        "db_write",
                        "success",
                        detail="仅生成标题",
                    )
            except CancelledError:
                summary["cancelled"] = True
                _emit_generation_event(on_event, task, "title", get_step(task, "title_request"), "cancelled", level="warning")
            except Exception as exc:
                summary["failed"] += 1
                summary["ok"] = False
                step = get_step(task, "title_request")
                error = _mark_generate_failed(task, exc, db_path, on_task_update)
                _emit_generation_event(on_event, task, "title", step, "failed", detail=error, level="error")
                _emit_generation_error(on_error, task, "title", step, exc, error)
            _emit_generation_progress(on_progress, summary)

    if not generate_cover:
        if summary["cancelled"]:
            summary["ok"] = False
        return summary

    cover_tasks = [
        task for task in eligible
        if task.id in title_results
    ]
    with ThreadPoolExecutor(
        max_workers=max(1, int(generation_cfg.get("image_concurrency", 1)))
    ) as executor:
        futures = {}
        for task in cover_tasks:
            if should_stop():
                summary["cancelled"] = True
                break
            new_title = title_results[task.id]
            try:
                set_step(task, "cover_prompt_render")
                _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "start")
                rendered_cover_prompt = prompt_module.render_prompt(
                    cover_prompt,
                    _prompt_context(task, new_title, account_by_alias),
                )
                _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "success")
                set_step(task, "cover_submit")
                _emit_generation_event(on_event, task, "cover", "cover_submit", "start")
                futures[
                    executor.submit(
                        gen_cover,
                        rendered_cover_prompt,
                        getattr(task, "old_cover_path", "") or "",
                        _new_cover_path(task, account_by_alias, image_root),
                        resolution=generation_cfg.get("resolution"),
                        jpg_quality=generation_cfg.get("jpg_quality"),
                        retry=generation_cfg.get("retry"),
                        config=config,
                        models_path=models_path,
                        on_step=step_callback(task, "cover"),
                        on_event=step_callback(task, "cover"),
                        cmhub_config_path=cmhub_config_path,
                    )
                ] = (task, new_title)
            except Exception as exc:
                summary["failed"] += 1
                summary["ok"] = False
                step = get_step(task, "cover_prompt_render")
                error = _mark_generate_failed(task, exc, db_path, on_task_update)
                _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error")
                _emit_generation_error(on_error, task, "cover", step, exc, error)
                _emit_generation_progress(on_progress, summary)
        for future in as_completed(futures):
            task, new_title = futures[future]
            if should_stop():
                summary["cancelled"] = True
                _cancel_pending(futures)
            try:
                new_cover_path = future.result()
                set_step(task, "db_write")
                _emit_generation_event(on_event, task, "cover", "db_write", "start")
                db.set_generated(task.id, new_title, new_cover_path, path=db_path)
                summary["cover_done"] += 1
                summary["generated_done"] += 1
                if on_task_update is not None:
                    on_task_update(
                        task.id,
                        {
                            "stage": "generated",
                            "status": "success",
                            "new_title": new_title,
                            "new_cover_path": new_cover_path,
                        },
                    )
                _emit_generation_event(on_event, task, "cover", "db_write", "success", detail=new_cover_path)
            except CancelledError:
                summary["cancelled"] = True
                _emit_generation_event(on_event, task, "cover", get_step(task, "cover_request"), "cancelled", level="warning")
            except Exception as exc:
                summary["failed"] += 1
                summary["ok"] = False
                step = get_step(task, "cover_request")
                error = _mark_generate_failed(task, exc, db_path, on_task_update)
                _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error")
                _emit_generation_error(on_error, task, "cover", step, exc, error)
            _emit_generation_progress(on_progress, summary)

    if summary["cancelled"]:
        summary["ok"] = False
    return summary


def fetch_cmhub_models(base_url, api_key, connect_timeout=10, read_timeout=30):
    """Fetch cmhub model aliases for settings UI."""

    base_url = str(base_url or "").strip()
    api_key = str(api_key or "")
    if not base_url:
        raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub Base URL")
    if not api_key:
        raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub API Key")
    data = _cmhub_call_with_retry(
        "GET",
        appconfig.cmhub_request_url(base_url, "/api/v1/models"),
        api_key,
        payload=None,
        connect_timeout=max(1, int(connect_timeout or 10)),
        read_timeout=max(1, int(read_timeout or 30)),
        attempts=1,
        on_retry=None,
    )
    models = data.get("models", [])
    if not isinstance(models, list):
        raise CMHubError("bad_response", "cmhub 模型列表格式错误")
    return [copy.deepcopy(model) for model in models if isinstance(model, dict)]


def fetch_cmhub_balance(base_url, api_key, connect_timeout=10, read_timeout=30):
    """Fetch cmhub point balance for settings UI."""

    base_url = str(base_url or "").strip()
    api_key = str(api_key or "")
    if not base_url:
        raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub Base URL")
    if not api_key:
        raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub API Key")
    data = _cmhub_call_with_retry(
        "GET",
        appconfig.cmhub_request_url(base_url, "/api/v1/balance"),
        api_key,
        payload=None,
        connect_timeout=max(1, int(connect_timeout or 10)),
        read_timeout=max(1, int(read_timeout or 30)),
        attempts=1,
        on_retry=None,
    )
    if "points_balance" not in data:
        raise CMHubError("bad_response", "cmhub 余额返回格式错误")
    return copy.deepcopy(data)


def _ai_backend(ai_cfg):
    backend = str(ai_cfg.get("backend", "direct") or "direct").strip().lower()
    if backend not in appconfig.AI_BACKENDS:
        raise AIError("AI backend 必须是 direct 或 cmhub")
    return backend


def _gen_title_cmhub(
    title_prompt,
    old_title,
    retry,
    config,
    cmhub_config_path,
    on_step=None,
    on_event=None,
):
    ai_cfg = appconfig.ai_config(config)
    _notify_step(on_step, "load_text_model")
    runtime = _cmhub_runtime(config, "title", cmhub_config_path)
    resolution = _normalize_cmhub_resolution(ai_cfg.get("resolution", "1k"))
    _notify_step(on_step, "title_build_request")
    payload = {
        "prompt": _cmhub_title_prompt(title_prompt, old_title),
        "model": runtime["alias"],
        "resolution": resolution,
    }
    attempts = _attempt_count(ai_cfg, retry)
    _notify_step(on_step, "title_request")
    data = _cmhub_call_with_retry(
        "POST",
        appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/title"),
        runtime["api_key"],
        payload=payload,
        connect_timeout=runtime["connect_timeout"],
        read_timeout=_cmhub_read_timeout(config, ai_cfg.get("resolution", "1k")),
        attempts=attempts,
        on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry(
            on_step,
            "title_request",
            attempt,
            total_attempts,
            exc,
        ),
    )
    _emit_cmhub_metadata(on_event, data, "title_request")
    _notify_step(on_step, "title_parse_response")
    titles = data.get("titles")
    if not isinstance(titles, list) or not titles:
        raise AIError("AI 返回为空标题")
    text = str(titles[0] or "").strip()
    if not text:
        raise AIError("AI 返回为空标题")
    return text


def _gen_cover_cmhub(
    cover_prompt,
    old_cover_path,
    out_path,
    resolution,
    jpg_quality,
    retry,
    config,
    cmhub_config_path,
    on_step=None,
    on_event=None,
):
    ai_cfg = appconfig.ai_config(config)
    _notify_step(on_step, "load_image_model")
    runtime = _cmhub_runtime(config, "image", cmhub_config_path)
    resolution = str(resolution or ai_cfg.get("resolution", "1k"))
    quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90))
    _notify_step(on_step, "cover_build_request")
    payload = {
        "prompt": str(cover_prompt or ""),
        "model": runtime["alias"],
        "image_base64": _image_data_url(old_cover_path),
        "resolution": _normalize_cmhub_resolution(resolution),
        "aspect_ratio": "1:1",
    }
    attempts = _attempt_count(ai_cfg, retry)
    read_timeout = _cmhub_read_timeout(config, resolution)
    _notify_step(on_step, "cover_request")
    data = _cmhub_call_with_retry(
        "POST",
        appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image"),
        runtime["api_key"],
        payload=payload,
        connect_timeout=runtime["connect_timeout"],
        read_timeout=read_timeout,
        attempts=attempts,
        on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry(
            on_step,
            "cover_request",
            attempt,
            total_attempts,
            exc,
        ),
    )
    _emit_cmhub_metadata(on_event, data, "cover_request")
    _notify_step(on_step, "cover_parse_response")
    image_url = str(data.get("image_url") or "").strip()
    if not image_url:
        raise AIError("AI 返回中没有图片数据")
    image_bytes = _download_cmhub_image(
        image_url,
        connect_timeout=runtime["connect_timeout"],
        read_timeout=read_timeout,
    )
    _notify_step(on_step, "cover_save")
    return _save_jpeg(image_bytes, out_path, resolution, quality)


def _cmhub_runtime(config, operation, cmhub_config_path):
    hub = appconfig.cmhub_config(config)
    api_key = appconfig.get_cmhub_api_key(path=cmhub_config_path)
    alias_key = "title_alias" if operation == "title" else "image_alias"
    missing = []
    if not hub.get("base_url"):
        missing.append("Base URL")
    if not api_key:
        missing.append("API Key")
    if not hub.get(alias_key):
        missing.append("生文别名" if operation == "title" else "生图别名")
    if missing:
        raise CMHubError(
            "cmhub_not_configured",
            "请去⑤设置配置 cmhub：缺少 " + "、".join(missing),
            retryable=False,
        )
    return {
        "base_url": hub["base_url"].rstrip("/"),
        "api_key": api_key,
        "alias": hub[alias_key],
        "connect_timeout": max(1, int(hub.get("connect_timeout", 10) or 10)),
    }


def _cmhub_title_prompt(title_prompt, old_title):
    return "%s\n\n旧标题：\n%s\n\n请只返回新标题，不要解释。" % (
        str(title_prompt or "").strip(),
        str(old_title or ""),
    )


def _normalize_cmhub_resolution(resolution):
    value = str(resolution or "1k").strip().lower()
    mapping = {
        "512": "512",
        "512x512": "512",
        "1k": "1K",
        "1K": "1K",
        "1024": "1K",
        "2k": "2K",
        "2K": "2K",
        "2048": "2K",
        "4k": "4K",
        "4K": "4K",
        "4096": "4K",
    }
    return mapping.get(value, str(resolution or "1K").upper())


def _cmhub_read_timeout(config, resolution):
    try:
        timeout = int(appconfig.response_timeout(config))
    except Exception:
        timeout = 600
    resolution_key = str(resolution or "").strip().lower()
    timeouts = appconfig.ai_config(config).get("resolution_timeouts", {})
    if resolution_key in timeouts:
        timeout = int(timeouts[resolution_key])
    return min(600, max(1, timeout))


def _cmhub_call_with_retry(
    method,
    url,
    api_key,
    payload,
    connect_timeout,
    read_timeout,
    attempts,
    on_retry=None,
):
    attempts = max(1, int(attempts or 1))
    last_exc = None
    for index in range(attempts):
        try:
            return _cmhub_call_once(
                method,
                url,
                api_key,
                payload,
                connect_timeout=connect_timeout,
                read_timeout=read_timeout,
            )
        except CMHubError as exc:
            last_exc = exc
            if not exc.retryable or index + 1 >= attempts:
                raise
            if on_retry is not None:
                try:
                    on_retry(index + 1, attempts, exc)
                except Exception:
                    pass
            time.sleep(_cmhub_retry_delay(exc, index))
    raise last_exc


def _cmhub_call_once(method, url, api_key, payload, connect_timeout, read_timeout):
    headers = {
        "Authorization": "Bearer " + str(api_key),
        "Accept": "application/json",
    }
    request_kwargs = {
        "headers": headers,
        "timeout": (max(1, int(connect_timeout)), max(1, int(read_timeout))),
    }
    if str(method).upper() != "GET":
        request_kwargs["json"] = payload or {}
    try:
        response = requests.request(str(method).upper(), url, **request_kwargs)
    except requests.exceptions.ConnectTimeout as exc:
        raise CMHubError("connect_timeout", "连接 cmhub 超时", retryable=True) from exc
    except requests.exceptions.ReadTimeout as exc:
        raise CMHubError("read_timeout", "等待 cmhub 返回超时", retryable=False) from exc
    except requests.exceptions.RequestException as exc:
        raise CMHubError(
            "network_error",
            _redact_cmhub(str(exc), api_key),
            retryable=False,
        ) from exc
    return _cmhub_response_json(response, api_key)


def _cmhub_response_json(response, api_key):
    try:
        data = response.json()
    except ValueError:
        data = {}
    status = getattr(response, "status_code", None)
    if status and status >= 400:
        raise _cmhub_error_from_response(data, response, api_key)
    if isinstance(data, dict) and isinstance(data.get("error"), dict):
        raise _cmhub_error_from_response(data, response, api_key)
    if not isinstance(data, dict):
        raise CMHubError("bad_response", "cmhub 返回格式错误", status=status)
    return data


def _cmhub_error_from_response(data, response, api_key):
    status = getattr(response, "status_code", None)
    error = data.get("error") if isinstance(data, dict) else None
    if not isinstance(error, dict):
        error = {}
    code = str(error.get("code") or _cmhub_code_for_status(status))
    raw_message = (error.get("message") or data.get("message")) if isinstance(data, dict) else ""
    if not raw_message:
        raw_message = getattr(response, "text", "")[:500]
    message = _cmhub_user_message(code, _redact_cmhub(raw_message or code, api_key))
    retry_after = _parse_retry_after(getattr(response, "headers", {}).get("Retry-After"))
    return CMHubError(
        code,
        message,
        status=status,
        retryable=_cmhub_retryable(code),
        retry_after=retry_after,
    )


def _cmhub_code_for_status(status):
    return {
        400: "bad_request",
        401: "unauthorized",
        402: "insufficient_points",
        403: "account_disabled",
        429: "rate_limited",
        502: "upstream_error",
    }.get(status, "unknown")


def _cmhub_retryable(code):
    return str(code) in {"upstream_error", "rate_limited", "connect_timeout"}


def _cmhub_user_message(code, message):
    defaults = {
        "insufficient_points": "点数不足，请先充值",
        "unauthorized": "cmhub API Key 无效，请去⑤设置重填",
        "account_disabled": "cmhub 账号已禁用，请去网页端处理",
        "bad_request": "cmhub 请求参数错误",
        "model_not_allowed": "cmhub 模型别名无权限",
        "no_pricing_rule": "cmhub 模型别名未配置价格",
        "content_blocked": "cmhub 内容安全策略拒绝本次生成",
        "upstream_error": "cmhub 上游生成失败，请稍后重试",
        "rate_limited": "cmhub 请求过于频繁，请稍后重试",
    }
    default = defaults.get(str(code))
    if default and message and str(message) not in default:
        return "%s：%s" % (default, message)
    return default or str(message or code)


def _parse_retry_after(value):
    try:
        if value is None or value == "":
            return None
        return max(0.0, float(value))
    except (TypeError, ValueError):
        return None


def _cmhub_retry_delay(exc, index):
    if exc.retry_after is not None:
        return min(2.0, max(0.1, float(exc.retry_after)))
    return min(2.0, 0.4 * (index + 1))


def _notify_cmhub_retry(callback, step, attempt, attempts, exc):
    if callback is None:
        return
    try:
        callback(
            {
                "step": step,
                "result": "retry",
                "attempt": attempt,
                "attempts": attempts,
                "detail": str(exc),
            }
        )
    except Exception:
        pass


def _emit_cmhub_metadata(callback, data, step):
    if callback is None or not isinstance(data, dict):
        return
    metadata = {
        key: data.get(key)
        for key in ("alias", "model_used", "points_cost", "points_balance", "call_id")
        if data.get(key) is not None
    }
    if not metadata:
        return
    try:
        callback(
            {
                "step": step,
                "result": "meta",
                "level": "info",
                "metadata": metadata,
            }
        )
    except Exception:
        pass


def _download_cmhub_image(url, connect_timeout, read_timeout, max_bytes=CMHUB_IMAGE_MAX_BYTES):
    _assert_public_http_url(url)
    try:
        response = requests.get(
            url,
            stream=True,
            timeout=(max(1, int(connect_timeout)), max(1, int(read_timeout))),
        )
    except requests.exceptions.RequestException as exc:
        raise AIError("下载 cmhub 图片失败: %s" % exc) from exc
    status = getattr(response, "status_code", 200)
    if status >= 400:
        raise AIError("下载 cmhub 图片失败: HTTP %s" % status)
    chunks = []
    total = 0
    iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
    for chunk in iterator:
        if not chunk:
            continue
        total += len(chunk)
        if total > max_bytes:
            raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
        chunks.append(chunk)
    return b"".join(chunks)


def _assert_public_http_url(url):
    parts = urllib.parse.urlsplit(str(url or ""))
    if parts.scheme not in {"http", "https"}:
        raise AIError("cmhub 图片地址只允许 http/https")
    host = parts.hostname
    if not host:
        raise AIError("cmhub 图片地址缺少域名")
    if _is_local_hostname(host):
        raise AIError("cmhub 图片地址不能指向本机或内网")
    try:
        _assert_public_ip(host)
        return
    except ValueError:
        pass
    try:
        addresses = socket.getaddrinfo(
            host,
            parts.port or (443 if parts.scheme == "https" else 80),
            type=socket.SOCK_STREAM,
        )
    except OSError as exc:
        raise AIError("cmhub 图片地址无法解析: %s" % exc) from exc
    if not addresses:
        raise AIError("cmhub 图片地址无法解析")
    for address in addresses:
        ip_text = address[4][0]
        _assert_public_ip(ip_text)


def _is_local_hostname(host):
    lowered = str(host or "").strip().lower().rstrip(".")
    return lowered in {"localhost"} or lowered.endswith(".localhost") or lowered.endswith(".local")


def _assert_public_ip(value):
    ip = ipaddress.ip_address(value)
    if (
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_multicast
        or ip.is_reserved
        or ip.is_unspecified
    ):
        raise AIError("cmhub 图片地址不能指向本机或内网")


def _redact_cmhub(text, api_key):
    return appconfig.redact_secrets(text, [api_key])

def _role_model(category, name, models_path):
    if not name:
        raise AIError("未配置默认 %s 模型" % category)
    model = appconfig.get_model(name, path=models_path)
    if model.get("category") != category:
        raise AIError("模型 %s 不是 %s 类别" % (name, category))
    if not model.get("enabled", True):
        raise AIError("模型已禁用: %s" % name)
    missing = [
        field
        for field in ("url", "model", "api_key")
        if not str(model.get(field, "")).strip()
    ]
    if missing:
        raise AIError("模型 %s 缺少字段: %s" % (name, ", ".join(missing)))
    return model


def _runtime_config(runtime):
    if runtime.get("config") is not None:
        return runtime["config"]
    config = appconfig.load_config()
    ai_updates = {
        key: value
        for key, value in runtime.items()
        if key in {
            "default_text_model",
            "default_image_model",
            "title_concurrency",
            "image_concurrency",
            "retry",
            "jpg_quality",
            "resolution",
            "generate_cover",
            "resolution_timeouts",
            "backend",
            "cmhub",
        }
    }
    if ai_updates:
        config = copy.deepcopy(config)
        config.setdefault("ai", {}).update(ai_updates)
    return config


def _prompt_value(prompt_values, name):
    if isinstance(prompt_values, dict):
        return str(
            prompt_values.get(name)
            or prompt_values.get(f"{name}_prompt")
            or ""
        )
    return ""


def _prompt_context(task, new_title, account_by_alias):
    return {
        "old_title": getattr(task, "old_title", ""),
        "new_title": new_title,
        "item_id": getattr(task, "item_id", ""),
        "account_name": _account_name(task, account_by_alias),
        "alias": getattr(task, "alias", ""),
    }


def _account_name(task, account_by_alias):
    alias = str(getattr(task, "alias", "") or "").strip()
    account = account_by_alias.get(alias)
    if account is not None:
        return getattr(account, "account_name", "") or alias
    return getattr(task, "account_name", "") or alias


def _new_cover_path(task, account_by_alias, image_root):
    alias = str(getattr(task, "alias", "") or "").strip()
    account = account_by_alias.get(alias)
    return image_paths.task_image_path(image_root, task, account, "new")


def _mark_generate_failed(task, exc, db_path, on_task_update):
    error = diagnostics.redact_log_text(str(exc) or exc.__class__.__name__)
    db.mark_failed(task.id, "generate", error, path=db_path)
    if on_task_update is not None:
        on_task_update(task.id, {"status": "failed", "last_error": error})
    return error


def _emit_generation_event(
    callback,
    task,
    phase,
    step,
    result,
    detail=None,
    level="info",
    attempt=None,
    attempts=None,
    metadata=None,
):
    if callback is None:
        return
    payload = {
        "task": task,
        "phase": phase,
        "step": str(step),
        "result": result,
        "level": level,
    }
    if detail is not None:
        payload["detail"] = diagnostics.redact_log_text(detail)
    if attempt is not None:
        payload["attempt"] = attempt
    if attempts is not None:
        payload["attempts"] = attempts
    if metadata is not None:
        payload["metadata"] = appconfig.sanitize_for_log(metadata)
    try:
        callback(payload)
    except Exception:
        return


def _emit_generation_error(callback, task, phase, step, exc, error):
    if callback is None:
        return
    payload = {
        "task": task,
        "phase": phase,
        "step": str(step),
        "error": diagnostics.redact_log_text(error),
        "exception": exc,
    }
    try:
        callback(payload)
    except Exception:
        return


def _cancel_pending(futures):
    for future in futures:
        if not future.done():
            future.cancel()


def _emit_generation_progress(on_progress, summary):
    if on_progress is None:
        return
    payload = dict(summary)
    try:
        on_progress(payload)
    except TypeError:
        on_progress(
            payload.get("title_done", 0),
            payload.get("cover_done", 0),
            payload.get("failed", 0),
        )


def _attempt_count(ai_cfg, retry):
    retry_count = ai_cfg.get("retry", 2) if retry is None else retry
    return max(1, int(retry_count) + 1)


def _read_timeout(model, config):
    return int(model.get("timeout_seconds") or appconfig.response_timeout(config))


def _connect_timeout(model):
    return int(model.get("connect_timeout_seconds") or 30)


def _headers(model, content_type):
    return {
        "Authorization": "Bearer " + model["api_key"],
        "Content-Type": content_type,
    }


def _call_with_retry(
    model,
    body,
    config,
    attempts,
    request_kind,
    content_type=None,
    on_retry=None,
):
    last_exc = None
    for index in range(attempts):
        try:
            return _call_once(model, body, config, request_kind, content_type=content_type)
        except Exception as exc:
            last_exc = exc
            if index + 1 >= attempts:
                break
            if on_retry is not None:
                try:
                    on_retry(index + 1, attempts, exc)
                except Exception:
                    pass
            time.sleep(min(2.0, 0.4 * (index + 1)))
    raise AIError(
        "AI 调用失败（已尝试 %s 次）: %s"
        % (attempts, _redact(str(last_exc), model))
    ) from last_exc


def _call_once(model, body, config, request_kind, content_type=None):
    if request_kind == "json":
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
        content_type = "application/json"
    else:
        data = body
    request = urllib.request.Request(
        appconfig.model_request_url(model),
        data=data,
        headers=_headers(model, content_type),
        method="POST",
    )
    timeout = max(_connect_timeout(model), _read_timeout(model, config))
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            raw = response.read()
    except urllib.error.HTTPError as exc:
        detail = _safe_http_error(exc, model)
        raise AIError("HTTP %s: %s" % (exc.code, detail)) from exc
    except urllib.error.URLError as exc:
        raise AIError(_redact(str(exc.reason), model)) from exc
    try:
        return json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise AIError("AI 返回不是有效 JSON") from exc


def _safe_http_error(exc, model):
    try:
        body = exc.read(2048).decode("utf-8", errors="replace")
    except Exception:
        body = ""
    return _redact(body or str(exc), model)


def _chat_payload(model, messages):
    payload = {
        "model": model["model"],
        "messages": messages,
    }
    payload.update(copy.deepcopy(model.get("extra_body", {})))
    return payload


def _image_chat_payload(model, cover_prompt, old_cover_path, resolution):
    prompt = "%s\n\n目标分辨率：%s。" % (str(cover_prompt or "").strip(), resolution)
    content = [
        {"type": "text", "text": prompt.strip()},
        {
            "type": "image_url",
            "image_url": {"url": _image_data_url(old_cover_path)},
        },
    ]
    return _chat_payload(model, [{"role": "user", "content": content}])


def _image_edit_body(model, cover_prompt, old_cover_path, resolution):
    fields = {
        "model": model["model"],
        "prompt": str(cover_prompt or ""),
        "size": _resolution_size_text(resolution),
    }
    fields.update(copy.deepcopy(model.get("extra_body", {})))
    files = {
        "image": (
            os.path.basename(old_cover_path),
            open(old_cover_path, "rb").read(),
            mimetypes.guess_type(old_cover_path)[0] or "application/octet-stream",
        )
    }
    return _multipart_body(fields, files)


def _multipart_body(fields, files):
    boundary = "----cmshopee-%s" % uuid.uuid4().hex
    chunks = []
    for name, value in fields.items():
        chunks.extend(
            [
                ("--%s\r\n" % boundary).encode("utf-8"),
                ('Content-Disposition: form-data; name="%s"\r\n\r\n' % name).encode("utf-8"),
                str(value).encode("utf-8"),
                b"\r\n",
            ]
        )
    for name, file_info in files.items():
        filename, data, content_type = file_info
        chunks.extend(
            [
                ("--%s\r\n" % boundary).encode("utf-8"),
                (
                    'Content-Disposition: form-data; name="%s"; filename="%s"\r\n'
                    % (name, filename)
                ).encode("utf-8"),
                ("Content-Type: %s\r\n\r\n" % content_type).encode("utf-8"),
                data,
                b"\r\n",
            ]
        )
    chunks.append(("--%s--\r\n" % boundary).encode("utf-8"))
    return b"".join(chunks), "multipart/form-data; boundary=%s" % boundary


def _image_data_url(path):
    mime = mimetypes.guess_type(path)[0] or "image/jpeg"
    with open(path, "rb") as fh:
        encoded = base64.b64encode(fh.read()).decode("ascii")
    return "data:%s;base64,%s" % (mime, encoded)


def _extract_text(data):
    if isinstance(data, dict):
        for key in ("output_text", "text", "content"):
            value = data.get(key)
            if isinstance(value, str):
                return value
        choices = data.get("choices")
        if isinstance(choices, list) and choices:
            first = choices[0]
            if isinstance(first, dict):
                if isinstance(first.get("text"), str):
                    return first["text"]
                message = first.get("message") or {}
                content = message.get("content")
                return _content_text(content)
    return ""


def _content_text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if not isinstance(item, dict):
                continue
            value = item.get("text") or item.get("content")
            if isinstance(value, str):
                parts.append(value)
        return "".join(parts)
    return ""


def _extract_image_bytes(data, model, config):
    image_ref = _find_image_ref(data)
    if not image_ref:
        raise AIError("AI 返回中没有图片数据")
    if image_ref.startswith("data:"):
        return _decode_data_url(image_ref)
    if _looks_base64(image_ref):
        return base64.b64decode(image_ref)
    return _download_image(image_ref, model, config)


def _find_image_ref(value):
    if isinstance(value, dict):
        for key in ("b64_json", "base64", "image_base64", "image", "url"):
            candidate = value.get(key)
            if isinstance(candidate, str) and candidate.strip():
                return candidate.strip()
        image_url = value.get("image_url")
        if isinstance(image_url, str):
            return image_url
        if isinstance(image_url, dict):
            candidate = image_url.get("url")
            if isinstance(candidate, str):
                return candidate
        for key in ("data", "choices", "output", "content", "images"):
            candidate = _find_image_ref(value.get(key))
            if candidate:
                return candidate
        message = value.get("message")
        if message is not None:
            candidate = _find_image_ref(message)
            if candidate:
                return candidate
    elif isinstance(value, list):
        for item in value:
            candidate = _find_image_ref(item)
            if candidate:
                return candidate
    return None


def _decode_data_url(value):
    if "," not in value:
        raise AIError("图片 data URL 格式错误")
    return base64.b64decode(value.split(",", 1)[1])


def _looks_base64(value):
    compact = value.strip()
    if compact.startswith(("http://", "https://")):
        return False
    if len(compact) < 32:
        return False
    allowed = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=\n\r")
    return all(char in allowed for char in compact)


def _download_image(url, model, config):
    request = urllib.request.Request(url, method="GET")
    try:
        with urllib.request.urlopen(request, timeout=_read_timeout(model, config)) as response:
            return response.read()
    except Exception as exc:
        raise AIError("下载 AI 图片失败: %s" % _redact(str(exc), model)) from exc


def _save_jpeg(image_bytes, out_path, resolution, jpg_quality):
    try:
        from PIL import Image
    except ImportError as exc:
        raise AIError("缺少 Pillow，无法保存 AI 图片") from exc

    import io

    out_path = os.path.abspath(str(out_path))
    directory = os.path.dirname(out_path)
    if directory:
        os.makedirs(directory, exist_ok=True)
    size = _resolution_size(resolution)
    try:
        with Image.open(io.BytesIO(image_bytes)) as image:
            image = image.convert("RGB")
            if size:
                image = image.resize(size, Image.LANCZOS)
            image.save(out_path, "JPEG", quality=jpg_quality, optimize=True)
    except Exception as exc:
        raise AIError("AI 图片保存失败: %s" % exc) from exc
    return out_path


def _resolution_size(resolution):
    return _RESOLUTION_SIZES.get(str(resolution))


def _resolution_size_text(resolution):
    size = _resolution_size(resolution)
    if not size:
        return str(resolution)
    return "%sx%s" % size


def _jpg_quality(value):
    value = int(value)
    return min(100, max(1, value))


def _redact(text, model):
    result = str(text)
    for secret in (model.get("api_key"),):
        if secret:
            result = result.replace(str(secret), "***")
    return result