"""数据库驱动的后台任务仓库与调度状态机。

已承载图片生成与批量水印；表结构刻意通用，后续可复用到 ZIP 导出等耗时操作。
Worker 不依赖 FastAPI 进程；第三方 provider task id 保存在子项中，重启后继续轮询。
"""
from __future__ import annotations

import uuid
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any

from sqlalchemy import case, func, or_, select, update

from db import get_session_factory
from models import BackgroundJob, BackgroundJobItem

JOB_QUEUED = "queued"
JOB_RUNNING = "running"
JOB_DONE = "done"
JOB_PARTIAL = "partial"
JOB_FAILED = "failed"
JOB_CANCELLED = "cancelled"

ITEM_QUEUED = "queued"
ITEM_RUNNING = "running"
ITEM_WAITING_PROVIDER = "waiting_provider"
ITEM_DONE = "done"
ITEM_FAILED = "failed"
ITEM_CANCELLED = "cancelled"

TERMINAL_ITEM_STATUSES = (ITEM_DONE, ITEM_FAILED, ITEM_CANCELLED)
TERMINAL_JOB_STATUSES = (JOB_DONE, JOB_PARTIAL, JOB_FAILED, JOB_CANCELLED)


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


@dataclass(frozen=True)
class ClaimedItem:
    job_id: str
    item_id: str
    job_kind: str
    product_id: str | None
    provider: str
    model: str | None
    platform: str
    lang: str
    ratio: str
    style_set: int
    request_payload: dict[str, Any]
    item_payload: dict[str, Any]
    provider_task_id: str | None
    provider_task_kind: str | None


async def create_image_job(
    *,
    product_id: str | None,
    kind: str,
    priority: int,
    provider: str,
    model: str | None,
    platform: str,
    lang: str,
    ratio: str,
    style_set: int,
    request_payload: dict[str, Any],
    item_payloads: list[dict[str, Any]],
) -> BackgroundJob:
    """一次事务写入父任务和逐图子项，API 成功返回即意味着任务可由 Worker 恢复。"""
    parsed_product_id = uuid.UUID(product_id) if product_id else None
    job = BackgroundJob(
        product_id=parsed_product_id,
        kind=kind,
        status=JOB_QUEUED,
        priority=priority,
        provider=provider,
        model=model,
        platform=platform,
        lang=lang,
        ratio=ratio,
        style_set=style_set,
        request_payload=request_payload,
        total=len(item_payloads),
    )
    async with get_session_factory()() as db:
        db.add(job)
        await db.flush()
        db.add_all([
            BackgroundJobItem(job_id=job.id, position=index, payload=payload, status=ITEM_QUEUED)
            for index, payload in enumerate(item_payloads)
        ])
        await db.commit()
        await db.refresh(job)
    return job


async def _refresh_job_status(db, job: BackgroundJob) -> None:
    """从子项计算父任务进度与终态；调用方已在同一事务中修改子项。"""
    rows = (await db.execute(
        select(BackgroundJobItem.status, func.count(BackgroundJobItem.id))
        .where(BackgroundJobItem.job_id == job.id)
        .group_by(BackgroundJobItem.status)
    )).all()
    counts = {status: int(count) for status, count in rows}
    done = counts.get(ITEM_DONE, 0)
    failed = counts.get(ITEM_FAILED, 0)
    cancelled = counts.get(ITEM_CANCELLED, 0)
    terminal = done + failed + cancelled
    job.completed = done
    job.failed = failed

    if terminal < job.total:
        job.status = JOB_RUNNING
        return
    job.completed_at = utcnow()
    if done == 0 and cancelled == job.total:
        job.status = JOB_CANCELLED
    elif failed == 0 and cancelled == 0:
        job.status = JOB_DONE
    elif done > 0:
        job.status = JOB_PARTIAL
    elif job.cancel_requested and cancelled > 0:
        job.status = JOB_CANCELLED
    else:
        job.status = JOB_FAILED

    errors = (await db.scalars(
        select(BackgroundJobItem.error)
        .where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status == ITEM_FAILED)
        .order_by(BackgroundJobItem.position)
    )).all()
    if errors:
        brief = [str(error)[:200] for error in errors if error]
        job.error = "；".join(brief[:6]) or "任务失败"


async def recover_expired_leases() -> int:
    """把上个 Worker 遗留的工作单元重新放回队列；provider id 不会被清空。"""
    now = utcnow()
    async with get_session_factory()() as db:
        result = await db.execute(
            update(BackgroundJobItem)
            .where(
                BackgroundJobItem.status == ITEM_RUNNING,
                BackgroundJobItem.lease_expires_at.is_not(None),
                BackgroundJobItem.lease_expires_at < now,
            )
            .values(status=ITEM_QUEUED, lease_owner=None, lease_expires_at=None, next_run_at=now)
        )
        await db.commit()
        return int(result.rowcount or 0)


async def claim_next_item(owner: str, lease_seconds: int = 90) -> ClaimedItem | None:
    """领取一个到期子项。单机 SQLite Worker 使用租约，未来迁移 PostgreSQL 后可直接强化为 SKIP LOCKED。"""
    now = utcnow()
    async with get_session_factory()() as db:
        # 先回收过期租约；即便 Worker 非正常退出，也不会永久卡住。
        await db.execute(
            update(BackgroundJobItem)
            .where(
                BackgroundJobItem.status == ITEM_RUNNING,
                BackgroundJobItem.lease_expires_at.is_not(None),
                BackgroundJobItem.lease_expires_at < now,
            )
            .values(status=ITEM_QUEUED, lease_owner=None, lease_expires_at=None, next_run_at=now)
        )
        query = (
            select(BackgroundJobItem, BackgroundJob)
            .join(BackgroundJob, BackgroundJob.id == BackgroundJobItem.job_id)
            .where(
                BackgroundJob.cancel_requested.is_(False),
                BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES),
                BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)),
                or_(BackgroundJobItem.next_run_at.is_(None), BackgroundJobItem.next_run_at <= now),
            )
            .order_by(BackgroundJob.priority.desc(), BackgroundJob.created_at.asc(), BackgroundJobItem.position.asc())
            .limit(1)
        )
        row = (await db.execute(query)).first()
        if row is None:
            await db.commit()
            return None
        item, job = row
        item.status = ITEM_RUNNING
        if not item.provider_task_id:
            item.attempts += 1
        item.lease_owner = owner
        item.lease_expires_at = now + timedelta(seconds=lease_seconds)
        item.next_run_at = None
        item.started_at = item.started_at or now
        job.status = JOB_RUNNING
        job.started_at = job.started_at or now
        await db.commit()
        return ClaimedItem(
            job_id=str(job.id), item_id=str(item.id), job_kind=job.kind,
            product_id=str(job.product_id) if job.product_id else None,
            provider=job.provider, model=job.model, platform=job.platform,
            lang=job.lang, ratio=job.ratio, style_set=job.style_set,
            request_payload=dict(job.request_payload or {}), item_payload=dict(item.payload or {}),
            provider_task_id=item.provider_task_id, provider_task_kind=item.provider_task_kind,
        )


async def mark_waiting_provider(item_id: str, provider_task_id: str, provider_task_kind: str, delay_seconds: int = 3) -> None:
    """持久化外部任务 ID 并安排下一次查询。

    取消可能恰好发生在 Worker 轮询供应商的网络请求期间。这里在同一事务内
    再次检查父任务，避免已取消任务被写回 waiting_provider 后永远不再被领取。
    """
    async with get_session_factory()() as db:
        item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
        if item is None:
            return
        job = await db.get(BackgroundJob, item.job_id)
        if job is None or job.cancel_requested:
            item.status = ITEM_CANCELLED
            item.lease_owner = None
            item.lease_expires_at = None
            item.next_run_at = None
            item.completed_at = utcnow()
            if job is not None:
                await _refresh_job_status(db, job)
            await db.commit()
            return
        item.provider_task_id = provider_task_id
        item.provider_task_kind = provider_task_kind
        item.status = ITEM_WAITING_PROVIDER
        item.lease_owner = None
        item.lease_expires_at = None
        item.next_run_at = utcnow() + timedelta(seconds=delay_seconds)
        await db.commit()


async def mark_item_done(
    item_id: str,
    *,
    result_url: str,
    source_url: str | None,
    watermark_config: dict | None,
    generated_asset_id: str | None,
) -> None:
    async with get_session_factory()() as db:
        item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
        if item is None:
            return
        job = await db.get(BackgroundJob, item.job_id)
        item.status = ITEM_DONE
        item.result_url = result_url
        item.source_url = source_url
        item.watermark_config = watermark_config
        item.generated_asset_id = uuid.UUID(generated_asset_id) if generated_asset_id else None
        item.error = None
        item.lease_owner = None
        item.lease_expires_at = None
        item.completed_at = utcnow()
        if job is not None:
            await _refresh_job_status(db, job)
        await db.commit()


async def mark_item_failed(item_id: str, error: str) -> None:
    async with get_session_factory()() as db:
        item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
        if item is None:
            return
        job = await db.get(BackgroundJob, item.job_id)
        item.status = ITEM_FAILED
        item.error = error[:500]
        item.lease_owner = None
        item.lease_expires_at = None
        item.completed_at = utcnow()
        if job is not None:
            await _refresh_job_status(db, job)
        await db.commit()


async def job_can_write(job_id: str) -> bool:
    """结果落盘前的最后一道确认：取消或商品被硬删除时禁止回写素材。"""
    async with get_session_factory()() as db:
        job = await db.get(BackgroundJob, uuid.UUID(job_id))
        if job is None or job.cancel_requested:
            return False
        if job.product_id is None:
            return True
        # 仅用主键查询避免导入 Product 形成循环；不存在即不允许生成结果复活商品。
        from models import Product
        return await db.get(Product, job.product_id) is not None


async def mark_item_cancelled(item_id: str) -> None:
    async with get_session_factory()() as db:
        item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
        if item is None:
            return
        job = await db.get(BackgroundJob, item.job_id)
        item.status = ITEM_CANCELLED
        item.lease_owner = None
        item.lease_expires_at = None
        item.completed_at = utcnow()
        if job is not None:
            await _refresh_job_status(db, job)
        await db.commit()


async def cancel_job(job_id: str) -> BackgroundJob | None:
    async with get_session_factory()() as db:
        job = await db.get(BackgroundJob, uuid.UUID(job_id))
        if job is None:
            return None
        if job.status in TERMINAL_JOB_STATUSES:
            return job
        job.cancel_requested = True
        await db.execute(
            update(BackgroundJobItem)
            .where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)))
            .values(status=ITEM_CANCELLED, completed_at=utcnow(), lease_owner=None, lease_expires_at=None)
        )
        await _refresh_job_status(db, job)
        await db.commit()
        await db.refresh(job)
        return job


async def delete_job(job_id: str) -> bool:
    """硬删除一个后台任务及其子项；调用方须先处理对应的存储对象。"""
    try:
        parsed = uuid.UUID(job_id)
    except ValueError:
        return False
    async with get_session_factory()() as db:
        job = await db.get(BackgroundJob, parsed)
        if job is None:
            return False
        await db.delete(job)
        await db.commit()
        return True


async def cancel_product_jobs(product_id: str) -> int:
    """硬删除商品前取消其未完成任务，Worker 随后只会丢弃结果而不会回写素材。"""
    try:
        parsed = uuid.UUID(product_id)
    except ValueError:
        return 0
    async with get_session_factory()() as db:
        jobs = (await db.scalars(
            select(BackgroundJob).where(
                BackgroundJob.product_id == parsed,
                BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES),
            )
        )).all()
        now = utcnow()
        for job in jobs:
            job.cancel_requested = True
            await db.execute(
                update(BackgroundJobItem)
                .where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)))
                .values(status=ITEM_CANCELLED, completed_at=now, lease_owner=None, lease_expires_at=None)
            )
            await _refresh_job_status(db, job)
        await db.commit()
        return len(jobs)


async def retry_failed_items(job_id: str) -> BackgroundJob | None:
    """仅把失败子项重新排队；外部失败任务不可复用，清空 provider id 后重新提交。"""
    async with get_session_factory()() as db:
        job = await db.get(BackgroundJob, uuid.UUID(job_id))
        if job is None:
            return None
        now = utcnow()
        result = await db.execute(
            update(BackgroundJobItem)
            .where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status == ITEM_FAILED)
            .values(
                status=ITEM_QUEUED, error=None, provider_task_id=None, provider_task_kind=None,
                result_url=None, source_url=None, watermark_config=None, generated_asset_id=None,
                next_run_at=now, lease_owner=None, lease_expires_at=None, completed_at=None,
            )
        )
        if not result.rowcount:
            return job
        job.cancel_requested = False
        job.status = JOB_QUEUED
        job.completed_at = None
        job.error = None
        await _refresh_job_status(db, job)
        await db.commit()
        await db.refresh(job)
        return job


async def retry_failed_item(job_id: str, item_id: str) -> BackgroundJob | None:
    """只重试一个失败子项，供批量水印等任务逐图补偿。"""
    try:
        parsed_job = uuid.UUID(job_id)
        parsed_item = uuid.UUID(item_id)
    except ValueError:
        return None
    async with get_session_factory()() as db:
        job = await db.get(BackgroundJob, parsed_job)
        item = await db.get(BackgroundJobItem, parsed_item)
        if job is None or item is None or item.job_id != job.id:
            return None
        if item.status != ITEM_FAILED:
            return job
        # 水印任务等待期间，素材可能被用户再次处理过。逐项重试时改为以
        # 最新 current 版本为基线，避免把已经更新的图片误判为永久失败。
        if job.kind == "watermark_batch":
            raw_asset_id = (item.payload or {}).get("asset_id")
            try:
                from models import ProductAsset
                asset = await db.get(ProductAsset, uuid.UUID(str(raw_asset_id)))
            except (TypeError, ValueError):
                asset = None
            if asset is not None:
                payload = dict(item.payload or {})
                payload["expected_current_storage_key"] = asset.current_storage_key
                payload["expected_current_url"] = asset.stored_url
                item.payload = payload
        item.status = ITEM_QUEUED
        item.error = None
        item.provider_task_id = None
        item.provider_task_kind = None
        item.result_url = None
        item.source_url = None
        item.watermark_config = None
        item.generated_asset_id = None
        item.next_run_at = utcnow()
        item.lease_owner = None
        item.lease_expires_at = None
        item.completed_at = None
        job.cancel_requested = False
        job.status = JOB_QUEUED
        job.completed_at = None
        job.error = None
        await _refresh_job_status(db, job)
        await db.commit()
        await db.refresh(job)
        return job


async def get_job(job_id: str) -> tuple[BackgroundJob, list[BackgroundJobItem]] | None:
    try:
        parsed = uuid.UUID(job_id)
    except ValueError:
        return None
    async with get_session_factory()() as db:
        job = await db.get(BackgroundJob, parsed)
        if job is None:
            return None
        items = (await db.scalars(
            select(BackgroundJobItem).where(BackgroundJobItem.job_id == job.id).order_by(BackgroundJobItem.position)
        )).all()
        return job, items


async def list_product_jobs(
    product_id: str,
    active_only: bool = False,
    limit: int = 20,
    kinds: tuple[str, ...] | None = None,
) -> list[BackgroundJob]:
    try:
        parsed = uuid.UUID(product_id)
    except ValueError:
        return []
    async with get_session_factory()() as db:
        query = select(BackgroundJob).where(BackgroundJob.product_id == parsed)
        if kinds:
            query = query.where(BackgroundJob.kind.in_(kinds))
        if active_only:
            query = query.where(BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES))
        return (await db.scalars(query.order_by(BackgroundJob.created_at.desc()).limit(limit))).all()


def suite_view(job: BackgroundJob, items: list[BackgroundJobItem]) -> dict[str, Any]:
    """兼容现有 /suites/{id} 契约，同时增加任务类型、取消状态与子项状态。"""
    return {
        "id": str(job.id), "status": job.status, "kind": job.kind,
        "style_set": job.style_set, "platform": job.platform, "lang": job.lang,
        "ratio": job.ratio, "provider": job.provider, "model": job.model,
        "total": job.total, "completed": job.completed, "failed": job.failed,
        "cancel_requested": job.cancel_requested, "error": job.error,
        "created_at": job.created_at, "started_at": job.started_at, "completed_at": job.completed_at,
        "images": [
            {
                "asset_id": str(item.generated_asset_id) if item.generated_asset_id else None,
                "type_id": (item.payload or {}).get("kind", "custom"),
                "item_id": str(item.id),
                "name": (item.payload or {}).get("title", "AI 生图"),
                "url": item.result_url or "", "status": item.status,
                "error": item.error, "completed_at": item.completed_at,
            }
            for item in items
        ],
    }


def video_job_view(job: BackgroundJob, items: list[BackgroundJobItem]) -> dict[str, Any]:
    """轮播视频的轻量任务视图；仍复用 BackgroundJob 的恢复/取消/重试状态机。"""
    payload = dict(job.request_payload or {})
    item = items[0] if items else None
    item_payload = dict(item.payload or {}) if item else {}
    return {
        "id": str(job.id), "kind": job.kind, "status": job.status,
        "total": job.total, "completed": job.completed, "failed": job.failed,
        "cancel_requested": job.cancel_requested, "preset": payload.get("preset", "wb_ozon"),
        "source_asset_ids": [str(source.get("asset_id")) for source in payload.get("sources", [])],
        "slide_seconds": payload.get("slide_seconds", 0),
        "duration_seconds": item_payload.get("duration_seconds"),
        "width": item_payload.get("width", 810), "height": item_payload.get("height", 1080),
        "file_size_bytes": item_payload.get("file_size_bytes"),
        "created_at": job.created_at, "started_at": job.started_at, "completed_at": job.completed_at,
        "result": None if item is None else {
            "item_id": str(item.id), "status": item.status, "url": item.result_url, "error": item.error,
        },
        "error": job.error,
    }