"""独立后台 Worker：按子项粒度执行持久化图片任务。"""
from __future__ import annotations

import asyncio
import logging
import socket
import uuid

from config import get_settings
from db import get_session_factory
from models import ProductAsset
from services.image_assets import get_image_asset_service
from services.generator import OutpaintCanvas, _image_size, _refs_for_job, advance_generation, prepare_outpaint_canvas
from services.jobs import (
    ClaimedItem,
    claim_next_item,
    job_can_write,
    mark_item_cancelled,
    mark_item_done,
    mark_item_failed,
    mark_waiting_provider,
    recover_expired_leases,
)
from services.prompts import build_context, build_prompt
from services.prompts.common import requirements_block
from services.storage import get_storage
from services.suite_service import append_generated_asset
from services.tasks import TaskImage
from services.watermark import apply_watermark
from services.video_renderer import VideoRenderCancelled, render_slideshow

log = logging.getLogger("background.worker")


def _worker_owner() -> str:
    return f"{socket.gethostname()}-{uuid.uuid4().hex[:10]}"


def _image_meta(data: bytes) -> tuple[str, str]:
    return (".png", "image/png") if data[:8] == b"\x89PNG\r\n\x1a\n" else (".jpg", "image/jpeg")


async def _store_image(job_id: str, position: int, data: bytes, version: str) -> str:
    suffix, content_type = _image_meta(data)
    return await get_storage().save_bytes(data, f"suites/{job_id}/{version}/{position}{suffix}", content_type)


async def _cleanup_temporary(urls: list[str]) -> None:
    """append 成功后由 image_assets 搬运并清理；取消时主动清理本次临时文件。"""
    if not urls:
        return
    try:
        from services.image_assets import get_image_asset_service
        await get_image_asset_service().delete_temporary_urls(urls)
    except Exception:  # noqa: BLE001
        log.exception("清理取消任务的临时图片失败")


async def _process_watermark_item(claim: ClaimedItem) -> None:
    """批量水印的一个素材子项：基于 original 重做 current，提交后回收旧水印对象。"""
    if not await job_can_write(claim.job_id):
        await mark_item_cancelled(claim.item_id)
        return
    payload = claim.item_payload
    raw_asset_id = payload.get("asset_id")
    try:
        asset_id = uuid.UUID(str(raw_asset_id))
    except (TypeError, ValueError) as exc:
        await mark_item_failed(claim.item_id, "水印任务缺少合法素材 ID")
        return

    obsolete_key: str | None = None
    result_url = ""
    source_url = ""
    watermark_config: dict | None = None
    try:
        async with get_session_factory()() as db:
            asset = await db.get(ProductAsset, asset_id)
            if asset is None:
                raise RuntimeError("素材已删除")
            if asset.type != "img":
                raise RuntimeError("仅支持图片素材")

            expected_key = payload.get("expected_current_storage_key")
            expected_url = payload.get("expected_current_url")
            if expected_key and asset.current_storage_key != expected_key:
                raise RuntimeError("素材在等待期间已被其他操作更新，请重新提交水印任务")
            if not expected_key and expected_url and asset.stored_url != expected_url:
                raise RuntimeError("素材在等待期间已被其他操作更新，请重新提交水印任务")

            watermark = claim.request_payload.get("watermark") or {}
            if not watermark.get("enabled"):
                raise RuntimeError("水印设置已关闭")
            image_assets = get_image_asset_service()
            update = await image_assets.apply_watermark(asset, watermark)
            asset.error = None
            result_url = asset.stored_url or ""
            source_url = asset.original_url or ""
            watermark_config = asset.watermark_config
            obsolete_key = update.obsolete_key
            await db.commit()

        # 数据库已切换到新版本后再删旧对象；删除失败不影响用户可用的新图片。
        if obsolete_key:
            try:
                await get_image_asset_service().delete_key(obsolete_key)
            except Exception:  # noqa: BLE001
                log.exception("水印已切换，但旧水印对象删除失败：%s", obsolete_key)
        await mark_item_done(
            claim.item_id, result_url=result_url, source_url=source_url or None,
            watermark_config=watermark_config, generated_asset_id=None,
        )
    except Exception as exc:  # noqa: BLE001
        log.exception("后台水印任务 %s/%s 执行失败", claim.job_id, claim.item_id)
        await mark_item_failed(claim.item_id, str(exc))


async def _process_slideshow_video_item(claim: ClaimedItem) -> None:
    """一个视频任务只有一个子项；输入 URL 在提交时冻结，避免后续选择状态影响成片。"""
    payload = claim.request_payload
    sources = list(payload.get("sources") or [])
    source_urls = [str(source.get("url") or "") for source in sources]
    if not 2 <= len(source_urls) <= 5 or any(not url for url in source_urls):
        await mark_item_failed(claim.item_id, "视频任务缺少 2–5 张可用图片")
        return
    try:
        async def should_cancel() -> bool:
            return not await job_can_write(claim.job_id)

        rendered = await render_slideshow(source_urls, should_cancel=should_cancel)
        if not await job_can_write(claim.job_id):
            await mark_item_cancelled(claim.item_id)
            return
        key = f"products/{claim.product_id}/videos/{claim.job_id}/result.mp4"
        result_url = await get_storage().save_bytes(rendered.data, key, "video/mp4")
        async with get_session_factory()() as db:
            from models import BackgroundJobItem
            item = await db.get(BackgroundJobItem, uuid.UUID(claim.item_id))
            if item is not None:
                item.payload = {
                    **dict(item.payload),
                    "duration_seconds": rendered.duration_seconds,
                    "width": rendered.width,
                    "height": rendered.height,
                    "file_size_bytes": len(rendered.data),
                }
                await db.commit()
        await mark_item_done(claim.item_id, result_url=result_url, source_url=None, watermark_config=None, generated_asset_id=None)
    except VideoRenderCancelled:
        await mark_item_cancelled(claim.item_id)
    except Exception as exc:  # noqa: BLE001
        log.exception("后台轮播视频任务 %s/%s 执行失败", claim.job_id, claim.item_id)
        await mark_item_failed(claim.item_id, str(exc))


def _outpaint_prompt(requirements: str | None) -> str:
    """仅扩图不复用套图提示词，避免默认风格、营销构图或文字污染边缘生成。"""
    requested = (requirements or "").strip()
    if not requested:
        requested = "Continue the existing surroundings naturally and seamlessly."
    return "\n\n".join([
        requirements_block(requested),
        "OUTPAINT ONLY: Image 1 contains an original photo placed on a 3:4 canvas. Generate pixels only "
        "in the transparent masked margins outside the original photo. Do not crop, move, resize, redraw, "
        "retouch, recolor, relight, or otherwise alter any pixel of the original photo.",
        "Do not add marketing layouts, product captions, overlays, props, or visible text unless the STRICT "
        "REQUIREMENTS explicitly ask for them. Make the new margins blend naturally with the original edges.",
        "FINAL CHECK: output exactly 3:4 and change only the newly expanded margins.",
    ])


async def process_item(claim: ClaimedItem) -> None:
    """推进一次子任务。外部模型未完成时只写 provider id 并退出，让出 Worker。"""
    if claim.job_kind == "watermark_batch":
        await _process_watermark_item(claim)
        return
    if claim.job_kind == "slideshow_video":
        await _process_slideshow_video_item(claim)
        return
    if claim.job_kind not in {"image_suite", "image_edit"}:
        await mark_item_failed(claim.item_id, f"Worker 暂不支持任务类型: {claim.job_kind}")
        return

    settings = get_settings()
    payload = claim.request_payload
    job_item = claim.item_payload
    provider = claim.provider or settings.image_provider
    model = claim.model or {
        "tongyi": settings.dashscope_model,
        "rightapi": settings.rightapi_image_model,
        "pcore": "gpt-image-2",
    }.get(provider, settings.ark_image_model)
    try:
        refs = _refs_for_job(list(payload.get("ref_images") or []), job_item)
        operation = payload.get("operation", "edit")
        target_size = _image_size(provider, claim.ratio, is_wan=provider == "tongyi" and model.lower().startswith("wan"), model=model)
        edit_mask: str | None = None
        outpaint: OutpaintCanvas | None = None
        if operation == "outpaint":
            if len(refs) != 1:
                raise RuntimeError("仅扩图必须且只能使用一张参考底图")
            prompt = _outpaint_prompt(payload.get("requirements"))
            # 首次提交时才构造图片和蒙版；之后轮询不再向供应商重复上传。
            if not claim.provider_task_id:
                outpaint = await prepare_outpaint_canvas(refs[0], target_size)
                refs = [outpaint.image_data_url]
                edit_mask = outpaint.mask_data_url
        else:
            context = build_context(payload.get("context") or {}, fallback_name="")
            prompt_extra = job_item.get("prompt_extra") or job_item
            prompt = build_prompt(
                provider, model, job_item.get("kind", "custom"), context, claim.style_set, claim.lang,
                extra=prompt_extra, style_prompt=payload.get("style_prompt"), requirements=payload.get("requirements"),
                requirements_exclusive=bool(payload.get("requirements_exclusive")),
            )
        is_wan = provider == "tongyi" and model.lower().startswith("wan")
        result = await advance_generation(
            provider, prompt, refs,
            size=target_size, model=model,
            provider_task_id=claim.provider_task_id, provider_task_kind=claim.provider_task_kind,
            edit_mask=edit_mask,
        )
        if result.data is None:
            if not result.provider_task_id or not result.provider_task_kind:
                raise RuntimeError("异步生成未返回可恢复的 provider task id")
            # 轮询供应商期间用户可能刚好取消任务。此时不能再写回
            # waiting_provider，否则父任务已取消、子项却会永久显示“进行中”。
            if not await job_can_write(claim.job_id):
                await mark_item_cancelled(claim.item_id)
                return
            await mark_waiting_provider(claim.item_id, result.provider_task_id, result.provider_task_kind)
            return

        # 生成完成后仍检查取消/商品删除，避免在用户删除商品后把素材回写回来。
        if not await job_can_write(claim.job_id):
            await mark_item_cancelled(claim.item_id)
            return

        if operation == "outpaint":
            # 异步任务完成时本轮未必创建过画布，按同一源图与尺寸确定性重建后覆盖原图区域。
            outpaint = outpaint or await prepare_outpaint_canvas(refs[0], target_size)
            result = type(result)(data=outpaint.restore_original(result.data))

        watermark = payload.get("watermark") or {}
        position = int(job_item.get("position", 0))
        if watermark.get("enabled"):
            source_url = await _store_image(claim.job_id, position, result.data, "original")
            result_url = await _store_image(claim.job_id, position, apply_watermark(result.data, watermark), "watermarked")
            watermark_config = watermark
        else:
            source_url = ""
            result_url = await _store_image(claim.job_id, position, result.data, "result")
            watermark_config = None

        asset_id: str | None = None
        if payload.get("append") and claim.product_id:
            image = TaskImage(
                type_id=job_item.get("kind", "custom"),
                name=job_item.get("title") or "AI 生图",
                url=result_url,
                source_url=source_url,
                watermark_config=watermark_config,
            )
            try:
                asset_id = await append_generated_asset(
                    claim.product_id, image, after_asset_id=payload.get("after_asset_id"),
                )
            except Exception:
                # 回写失败时保留 suites 临时文件和失败状态，避免用户看不到实际结果而又无法诊断。
                raise
            result_url, source_url, watermark_config = image.url, image.source_url, image.watermark_config

        await mark_item_done(
            claim.item_id, result_url=result_url, source_url=source_url or None,
            watermark_config=watermark_config, generated_asset_id=asset_id,
        )
    except Exception as exc:  # noqa: BLE001
        log.exception("后台图片任务 %s/%s 执行失败", claim.job_id, claim.item_id)
        await mark_item_failed(claim.item_id, str(exc))


async def run_worker(*, once: bool = False, idle_seconds: float = 0.75) -> None:
    owner = _worker_owner()
    recovered = await recover_expired_leases()
    if recovered:
        log.warning("恢复了 %d 个过期后台任务租约", recovered)
    log.info("后台 Worker 已启动：%s", owner)
    while True:
        claim = await claim_next_item(owner, lease_seconds=get_settings().background_job_lease_seconds)
        if claim is None:
            if once:
                return
            await asyncio.sleep(idle_seconds)
            continue
        await process_item(claim)
        if once:
            return