"""AI Labs scheduled jobs: the generation reaper and the retention sweep.

**``ai_labs.reap_generations``** is what makes a generation durable rather
than merely long-running. A job is started by the request that created it,
in the web process; a deploy, a crash or an OOM in the middle of a
90-second video would otherwise leave the row ``running`` forever with an
employee watching a spinner. The reaper looks for unfinished jobs whose
lease has expired and starts them again — and because the vendor task id
was persisted the moment the vendor accepted it, "again" means *resume
polling the work already in flight*, not submit a second render.

It **spawns** rather than awaits. The reaper's responsibility is to restart
stalled work, not to sit through it; a handler that awaited a 15-minute
video deadline would hold a pool connection for 15 minutes and make the
next tick late. The lease is what keeps a spawned job from being started
twice: a job with a live heartbeat is not reapable.

**``ai_labs.sweep_retention``** is a no-op unless an operator sets
``ai_labs.retention_days``. The default is to keep an employee's own
workspace until they delete it — an unannounced expiry on someone's saved
work is a surprise, and surprises about deletion are the worst kind.
"""
from __future__ import annotations

import logging
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy import delete, select
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.background import spawn
from app.core.storage import delete_object
from app.core.storage_uploads import forget_upload_grants
from app.domains.ai_labs import executor
from app.domains.ai_labs.models import (
    GEN_FAILED,
    GEN_QUEUED,
    GEN_RUNNING,
    AiLabConversation,
    AiLabGeneration,
    AiLabMessage,
)
from app.domains.ai_labs.settings import resolve_config
from app.scheduled_jobs.registry import register_scheduled_job

logger = logging.getLogger(__name__)

#: Jobs restarted per tick. A burst that outlived a deploy drains over a
#: few ticks rather than opening dozens of vendor pollers at once.
MAX_RESUMED_PER_TICK = 10

#: Rows deleted per retention batch, committed per batch. Long
#: transactions on ``identity_persons``-referencing tables are how this
#: platform previously FK-locked the person table and browned out login.
RETENTION_BATCH = 100


@register_scheduled_job(
    key="ai_labs.reap_generations",
    display_name="AI 实验室 生成任务续跑",
    description=(
        "找出租约已过期但仍未完成的生成任务并继续驱动它们——后端重启或崩溃"
        "打断的视频/图片任务由此恢复(厂商任务 id 已落库,故是继续轮询而非"
        "重新提交);重试次数用尽的任务在此判失败,不会永久占着队列。"
    ),
)
async def run_reap_generations(
    db: AsyncSession, config: dict[str, Any],
) -> dict[str, Any]:
    """Resume orphaned generation jobs; fail the ones out of attempts."""
    settings_config = await resolve_config(db)
    if not settings_config.available:
        # Nothing can be driven without a key, and failing every queued
        # job because a connector is briefly missing would be worse than
        # leaving them to the next tick.
        return {"skipped": "ai_labs_unavailable"}

    try:
        candidates = await executor.find_reapable(
            db, limit=MAX_RESUMED_PER_TICK * 3,
        )
    except OperationalError:
        # A deadlock or lost connection poisons the session; surface it so
        # the scheduler's transient-infra retry sees it rather than
        # recording a green run.
        logger.exception("ai_labs: reaper could not list reapable jobs")
        raise

    exhausted = 0
    resumed = 0
    now = datetime.now(UTC)

    for generation_id in candidates:
        row = await db.get(AiLabGeneration, generation_id)
        if row is None or row.status not in (GEN_QUEUED, GEN_RUNNING):
            continue
        if row.attempts >= executor.MAX_ATTEMPTS:
            # Out of attempts. Failing it here — rather than letting the
            # reaper pick it up forever — is what stops one poisonous job
            # from consuming a slot on every tick.
            row.status = GEN_FAILED
            row.finished_at = now
            row.lease_owner = None
            row.lease_expires_at = None
            # ``_release_for_retry`` left the retry sentinel on the row, so
            # "keep whatever code is there" would preserve ``RETRYING`` on a
            # job that is never going to be retried again — a terminal row
            # claiming to be mid-retry. Replace the sentinel; keep a real
            # vendor code if one was recorded.
            if row.error_code in (None, "", executor.RETRY_SENTINEL):
                row.error_code = "AI_LABS_ATTEMPTS_EXHAUSTED"
            row.error_message = (
                "This job failed repeatedly and was not retried again. "
                "Please submit it once more."
            )
            exhausted += 1
            continue
        if resumed >= MAX_RESUMED_PER_TICK:
            continue
        spawn(
            executor.run_job(generation_id),
            name=f"ai_labs.reap.{generation_id}",
        )
        resumed += 1

    await db.commit()
    return {
        "candidates": len(candidates),
        "resumed": resumed,
        "failed_exhausted": exhausted,
    }


@register_scheduled_job(
    key="ai_labs.sweep_retention",
    display_name="AI 实验室 历史保留期清理",
    description=(
        "当 ai_labs.retention_days > 0 时,删除超过保留期的对话与作品(含 S3 "
        "文件)。默认 0 = 不清理:AI 实验室是员工自己的工作区,未经宣告的"
        "过期删除是意外,而删除类的意外是最糟的一种。"
    ),
)
async def run_sweep_retention(
    db: AsyncSession, config: dict[str, Any],
) -> dict[str, Any]:
    """Delete history past the configured retention window."""
    settings_config = await resolve_config(db)
    days = settings_config.retention_days
    if days <= 0:
        return {"skipped": "retention_disabled"}

    cutoff = datetime.now(UTC) - timedelta(days=days)
    deleted_conversations = 0
    deleted_generations = 0
    objects: list[str] = []

    # ── conversations ──
    while True:
        ids = list(
            (
                await db.execute(
                    select(AiLabConversation.id)
                    .where(AiLabConversation.created_at < cutoff)
                    .limit(RETENTION_BATCH),
                )
            ).scalars().all(),
        )
        if not ids:
            break
        attachment_rows = (
            await db.execute(
                select(AiLabMessage.attachments).where(
                    AiLabMessage.conversation_id.in_(ids),
                ),
            )
        ).scalars().all()
        for attachments in attachment_rows:
            for item in attachments or []:
                key = (item or {}).get("object_key")
                if isinstance(key, str) and key:
                    objects.append(key)
        await db.execute(
            delete(AiLabMessage).where(
                AiLabMessage.conversation_id.in_(ids),
            ),
        )
        await db.execute(
            delete(AiLabConversation).where(AiLabConversation.id.in_(ids)),
        )
        await db.commit()
        deleted_conversations += len(ids)
        if len(ids) < RETENTION_BATCH:
            break

    # ── generations ──
    while True:
        rows = (
            await db.execute(
                select(
                    AiLabGeneration.id,
                    AiLabGeneration.artifacts,
                    AiLabGeneration.input_object_keys,
                )
                .where(AiLabGeneration.created_at < cutoff)
                .limit(RETENTION_BATCH),
            )
        ).all()
        if not rows:
            break
        ids = [row[0] for row in rows]
        for _id, artifacts, inputs in rows:
            for item in artifacts or []:
                key = (item or {}).get("object_key")
                if isinstance(key, str) and key:
                    objects.append(key)
            for key in inputs or []:
                if isinstance(key, str) and key:
                    objects.append(key)
        await db.execute(
            delete(AiLabGeneration).where(AiLabGeneration.id.in_(ids)),
        )
        await db.commit()
        deleted_generations += len(ids)
        if len(rows) < RETENTION_BATCH:
            break

    # The grant ledger rows go with the content, for the same reason the
    # employee-initiated delete removes them: a ``referenced`` row keeps a
    # full content fingerprint of a file this job has just deleted. No
    # owner scoping here — this is a system sweep acting on rows it has
    # itself already removed, spanning every employee in the window.
    if objects:
        released = await forget_upload_grants(
            db, object_keys=objects, person_id=None,
        )
        await db.commit()
    else:
        released = 0

    # S3 last, after the rows are gone. A bucket failure must not leave
    # the database rows behind — the employee's data being deleted is the
    # commitment.
    #
    # **A stranded object here is stranded for good, and that is worth
    # saying plainly rather than deferring to a sweep that cannot see it.**
    # ``storage.orphan_upload_cleanup`` walks the ``storage_upload_grants``
    # ledger; ``ai_labs_output/`` keys were never granted (the executor
    # writes them server-side), and the input keys' grants are released by
    # this job below, which removes them from that ledger too. So the count
    # this job returns is the only signal, and a non-zero
    # ``objects_stranded`` is an operator's cue to remove them by hand.
    forgotten = 0
    stranded: list[str] = []
    for key in objects:
        try:
            await delete_object(db, object_key=key)
            forgotten += 1
        except Exception as exc:  # noqa: BLE001
            stranded.append(key)
            logger.warning(
                "ai_labs: retention could not delete object %s: %s", key, exc,
            )
    if stranded:
        logger.error(
            "ai_labs: retention left %s object(s) in the bucket with no row "
            "referencing them and no sweep that can find them; remove them "
            "by hand: %s",
            len(stranded), stranded[:20],
        )

    return {
        "retention_days": days,
        "conversations_deleted": deleted_conversations,
        "generations_deleted": deleted_generations,
        "objects_deleted": forgotten,
        "objects_seen": len(objects),
        "objects_stranded": len(stranded),
        "grants_released": released,
    }


__all__ = ["run_reap_generations", "run_sweep_retention"]