from __future__ import annotations

import json
import logging
from datetime import datetime, timedelta, timezone
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.enums.generation_task import (
    ALLOWED_GENERATION_MODES,
    ChatGenerationPipelineStage,
    ChatGenerationTaskEventType,
    ChatGenerationTaskStatus,
)
from app.models.chat_generation_task import ChatGenerationTask
from app.services.celery_download_recovery_service import (
    ensure_aware_utc,
    get_download_active_payloads,
    get_due_download_record_ids,
    postpone_download_active_check,
    remove_download_active,
)
from app.services.generation_log_service import log_task_event
from app.services.generation_module_hook_service import notify_chat_generation_task_finished
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
from app.services.redis_registry_service import (
    redis_get_due_registry_ids,
    redis_get_registry_payloads,
    redis_postpone_registry_item,
    redis_remove_registry_item,
)

logger = logging.getLogger("video_gen")

POLL_QUEUE = "gen_provider_poll"


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


def _is_expired(value: datetime | None, now: datetime | None = None) -> bool:
    checked = ensure_aware_utc(value)
    if checked is None:
        return True
    return checked <= (now or _now())


def _queue_timeout_at(task: ChatGenerationTask, now: datetime | None = None) -> datetime:
    current_time = now or _now()
    enqueued_at = ensure_aware_utc(task.download_enqueued_at)
    if enqueued_at is None:
        return current_time
    return enqueued_at + timedelta(seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))


def _is_queue_timeout(task: ChatGenerationTask, now: datetime | None = None) -> bool:
    current_time = now or _now()
    return _queue_timeout_at(task, current_time) <= current_time


def _is_final_task_state(task: ChatGenerationTask) -> bool:
    return task.status in (ChatGenerationTaskStatus.COMPLETED.value, ChatGenerationTaskStatus.FAILED.value) or task.pipeline_stage in (
        ChatGenerationPipelineStage.DONE.value,
        ChatGenerationPipelineStage.FAILED.value,
        ChatGenerationPipelineStage.TIMEOUT.value,
        ChatGenerationPipelineStage.DOWNLOAD_FAILED.value,
    )


def _is_success(status: str | None) -> bool:
    return str(status or "").lower() in ("succeeded", "success", "completed", "done")


def _is_failed(status: str | None) -> bool:
    return str(status or "").lower() in ("failed", "error", "canceled", "cancelled")


def _engine_snapshot(task: ChatGenerationTask) -> dict[str, Any]:
    try:
        value = json.loads(task.engine_snapshot_json or "{}")
        return value if isinstance(value, dict) else {}
    except Exception:
        return {}


def _poll_queue_timeout_at(now: datetime | None = None) -> datetime:
    current_time = now or _now()
    return current_time + timedelta(seconds=int(settings.POLL_TASK_QUEUE_TIMEOUT_SECONDS or 120))


async def _remove_poll_active(task_id: str) -> None:
    await redis_remove_registry_item(
        hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
        item_id=task_id,
        log_context="poll_active",
    )


async def _postpone_poll_active(
    *,
    task_id: str,
    payload: dict[str, Any] | None = None,
    check_at: datetime | int | float | None = None,
) -> None:
    await redis_postpone_registry_item(
        hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
        item_id=task_id,
        payload=payload,
        check_at=check_at or _poll_queue_timeout_at(),
        log_context="poll_active",
    )


async def recover_one_download_task(
    db: AsyncSession,
    task: ChatGenerationTask,
    *,
    payload: dict[str, Any] | None = None,
    source: str = "startup_db",
) -> str:
    from app.tasks.generation_download_tasks import (
        DOWNLOAD_STAGE_DOWNLOADING,
        DOWNLOAD_STAGE_QUEUED,
        DOWNLOAD_STAGE_RETRY_WAITING,
        enqueue_download_task,
    )

    current_time = _now()

    if not task:
        return "skip_missing_task"
    if task.generation_mode not in ALLOWED_GENERATION_MODES:
        await remove_download_active(task.id)
        return "clean_invalid_mode"
    if _is_final_task_state(task):
        await remove_download_active(task.id)
        return "clean_final_state"
    if task.status != ChatGenerationTaskStatus.GENERATING.value:
        await remove_download_active(task.id)
        await log_task_event(
            task,
            event_type=ChatGenerationTaskEventType.DOWNLOAD_SKIP_NOT_GENERATING.value,
            message=f"{source} 下载恢复跳过：任务不是 generating",
            detail={"status": task.status, "stage": task.pipeline_stage},
        )
        return "clean_not_generating"
    if not task.remote_result_url:
        await log_task_event(
            task,
            event_type=ChatGenerationTaskEventType.DOWNLOAD_SKIP_NO_REMOTE_RESULT_URL.value,
            message=f"{source} 下载恢复跳过：缺少 remote_result_url",
            detail={"status": task.status, "stage": task.pipeline_stage},
        )
        return "skip_no_remote_result_url"

    stage = task.pipeline_stage
    redis_payload = payload or {}

    if stage == ChatGenerationPipelineStage.RESULT_READY.value:
        await log_task_event(
            task,
            event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
            message=f"{source} 发现 result_ready 未完成下载，启动时恢复投递下载任务",
            detail={"payload": redis_payload},
        )
        await enqueue_download_task(
            db,
            task,
            recover=True,
            reason=f"{source}_result_ready",
        )
        return "recover_result_ready"

    if stage == DOWNLOAD_STAGE_QUEUED:
        if _is_queue_timeout(task, current_time):
            await log_task_event(
                task,
                event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
                message=f"{source} 发现 download_queued 长时间未消费，启动时恢复投递下载任务",
                detail={"payload": redis_payload},
            )
            await enqueue_download_task(
                db,
                task,
                recover=True,
                reason=f"{source}_download_queued_timeout",
            )
            return "recover_queued_timeout"

        await postpone_download_active_check(
            record_id=task.id,
            payload=payload,
            check_at=_queue_timeout_at(task, current_time),
        )
        return "skip_queued_not_timeout"

    if stage == DOWNLOAD_STAGE_DOWNLOADING:
        if _is_expired(task.download_lease_until, current_time):
            await log_task_event(
                task,
                event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
                message=f"{source} 发现 downloading lease 过期，启动时恢复投递下载任务",
                detail={"payload": redis_payload},
            )
            await enqueue_download_task(
                db,
                task,
                recover=True,
                reason=f"{source}_downloading_lease_expired",
            )
            return "recover_downloading_expired"

        await postpone_download_active_check(
            record_id=task.id,
            payload=payload,
            check_at=task.download_lease_until,
        )
        return "skip_downloading_alive"

    if stage == DOWNLOAD_STAGE_RETRY_WAITING:
        if _is_expired(task.download_next_retry_at, current_time):
            await log_task_event(
                task,
                event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
                message=f"{source} 发现 retry_waiting 到期，启动时恢复投递下载任务",
                detail={"payload": redis_payload},
            )
            await enqueue_download_task(
                db,
                task,
                recover=True,
                reason=f"{source}_retry_waiting_due",
            )
            return "recover_retry_due"

        await postpone_download_active_check(
            record_id=task.id,
            payload=payload,
            check_at=task.download_next_retry_at,
        )
        return "skip_retry_waiting_not_due"

    return f"skip_stage_{stage}"


async def recover_download_tasks_once(db: AsyncSession) -> dict[str, Any]:
    """启动时下载容灾扫描。

    先按 Redis active_index 找到到期下载任务；Redis 不可用或索引丢失时，
    再通过 DB fallback 扫描 result_ready/download_* 状态，避免任务永久卡住。
    """
    checked_ids: set[str] = set()
    results: dict[str, int] = {}

    due_ids = await get_due_download_record_ids(
        limit=settings.DOWNLOAD_RECOVERY_BATCH_SIZE,
    )
    payloads = await get_download_active_payloads(due_ids)

    for task_id in due_ids:
        result = await db.execute(
            select(ChatGenerationTask)
            .where(
                ChatGenerationTask.id == task_id,
                ChatGenerationTask.deleted_at.is_(None),
            )
            .with_for_update()
            .limit(1)
        )
        task = result.scalar_one_or_none()
        if task is None:
            await remove_download_active(task_id)
            action = "clean_missing_task"
        else:
            checked_ids.add(task.id)
            action = await recover_one_download_task(
                db,
                task,
                payload=payloads.get(task_id),
                source="startup_redis",
            )
        results[action] = results.get(action, 0) + 1

    # DB fallback：不依赖 Redis active 注册表。
    fallback_result = await db.execute(
        select(ChatGenerationTask)
        .where(
            ChatGenerationTask.deleted_at.is_(None),
            ChatGenerationTask.generation_mode.in_(list(ALLOWED_GENERATION_MODES)),
            ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
            ChatGenerationTask.remote_result_url.is_not(None),
            ChatGenerationTask.pipeline_stage.in_(
                [
                    ChatGenerationPipelineStage.RESULT_READY.value,
                    ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value,
                    ChatGenerationPipelineStage.DOWNLOADING.value,
                    ChatGenerationPipelineStage.RETRY_WAITING.value,
                ]
            ),
        )
        .order_by(ChatGenerationTask.updated_at.asc())
        .limit(int(settings.DOWNLOAD_RECOVERY_BATCH_SIZE or 100))
        .with_for_update(skip_locked=True)
    )
    fallback_tasks = fallback_result.scalars().all()

    for task in fallback_tasks:
        if task.id in checked_ids:
            continue
        action = await recover_one_download_task(
            db,
            task,
            payload=None,
            source="startup_db",
        )
        results[action] = results.get(action, 0) + 1
        checked_ids.add(task.id)

    return {"checked": len(checked_ids), "results": results}


async def _mark_timeout(
    db: AsyncSession,
    task: ChatGenerationTask,
    *,
    error_message: str = "任务超时",
) -> str:
    await mark_chat_generation_task_failed_and_refund_once(
        db,
        task=task,
        error_message=error_message,
        pipeline_stage=ChatGenerationPipelineStage.TIMEOUT.value,
    )
    await notify_chat_generation_task_finished(db, task)
    await db.commit()
    await _remove_poll_active(task.id)
    await log_task_event(
        task,
        event_type="TASK_TIMEOUT",
        to_status="failed",
        to_stage=ChatGenerationPipelineStage.TIMEOUT.value,
    )
    return "mark_timeout"


async def _mark_failed(
    db: AsyncSession,
    task: ChatGenerationTask,
    *,
    error_message: str,
    event_type: str = "POLL_FAILED",
    detail: Any = None,
) -> str:
    await mark_chat_generation_task_failed_and_refund_once(
        db,
        task=task,
        error_message=error_message,
        pipeline_stage=ChatGenerationPipelineStage.FAILED.value,
    )
    await notify_chat_generation_task_finished(db, task)
    await db.commit()
    await _remove_poll_active(task.id)
    await log_task_event(task, event_type=event_type, message=task.error_message, detail=detail)
    return "mark_failed"


async def recover_one_generation_task(
    db: AsyncSession,
    task: ChatGenerationTask,
    *,
    payload: dict[str, Any] | None = None,
    source: str = "startup_db",
) -> str:
    """恢复单个生成任务。

    分流原则：
    1. 已有 remote_result_url：只恢复下载，不 poll，不重新 create。
    2. 已有 provider_task_id/seedance_task_id：恢复 poll。
    3. 无结果 URL、无供应商任务 ID：deadline 未过才恢复 create。
    4. 无结果 URL、无供应商任务 ID：deadline 已过直接超时失败，不再补救生成。
    """
    from app.tasks.generation_create_tasks import chatapi_create_generation_task
    from app.tasks.generation_download_tasks import enqueue_download_task
    from app.tasks.generation_poll_tasks import poll_generation_task, register_poll_active

    current_time = _now()
    redis_payload = payload or {}

    if not task:
        return "skip_missing_task"
    if task.generation_mode not in ALLOWED_GENERATION_MODES:
        await _remove_poll_active(task.id)
        return "clean_invalid_mode"
    if _is_final_task_state(task):
        await _remove_poll_active(task.id)
        return "clean_final_state"
    if task.status != ChatGenerationTaskStatus.GENERATING.value:
        await _remove_poll_active(task.id)
        return "clean_not_generating"

    has_remote_result = bool(str(task.remote_result_url or "").strip())
    has_provider_task_id = bool(str(task.provider_task_id or "").strip() or str(task.seedance_task_id or "").strip())
    is_deadline_expired = bool(task.deadline_at and _is_expired(task.deadline_at, current_time))

    # 最高优先级：只要远程结果 URL 已经落库，说明生成侧已经成功。
    # 不管当前 pipeline_stage 是 queued/creating/waiting/result_ready/download_*，恢复时都不能重复 create 或 poll。
    if has_remote_result:
        await _remove_poll_active(task.id)
        await log_task_event(
            task,
            event_type="GENERATION_RECOVERY_ENQUEUE",
            message=f"{source} 发现任务已存在 remote_result_url，恢复投递下载队列",
            detail={
                "pipeline_stage": task.pipeline_stage,
                "payload": redis_payload,
                "deadline_expired": is_deadline_expired,
            },
        )
        await enqueue_download_task(
            db,
            task,
            recover=True,
            reason=f"{source}_has_remote_result_url",
        )
        return "recover_download_has_remote_result"

    # 已经过 deadline 且没有结果 URL：
    # - 有供应商任务 ID：交给 poll worker 做最后一次状态确认；
    # - 没有供应商任务 ID：说明没有可查询的远程任务，直接按超时失败处理，不再重新 create。
    if is_deadline_expired:
        if has_provider_task_id:
            task.pipeline_stage = ChatGenerationPipelineStage.WAITING_REMOTE.value
            await db.commit()
            await log_task_event(
                task,
                event_type="GENERATION_RECOVERY_ENQUEUE",
                message=f"{source} 发现任务已到 deadline 且存在供应商任务ID，投递 poll 队列做最终查询",
                detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
            )
            poll_generation_task.apply_async(args=[task.id], queue=POLL_QUEUE, countdown=0)
            await register_poll_active(
                task,
                check_at=_poll_queue_timeout_at(),
                reason=f"{source}_deadline_final_poll",
            )
            return "recover_deadline_final_poll"

        await log_task_event(
            task,
            event_type="GENERATION_RECOVERY_TIMEOUT",
            message=f"{source} 发现任务已到 deadline，且没有 remote_result_url/供应商任务ID，按超时失败处理",
            detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
        )
        return await _mark_timeout(db, task)

    # 未过 deadline：有供应商任务 ID 才允许恢复到 poll 队列。
    if has_provider_task_id:
        task.pipeline_stage = ChatGenerationPipelineStage.WAITING_REMOTE.value
        await db.commit()
        await log_task_event(
            task,
            event_type="GENERATION_RECOVERY_ENQUEUE",
            message=f"{source} 发现任务存在供应商任务ID，恢复投递轮询队列",
            detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
        )
        poll_generation_task.apply_async(args=[task.id], queue=POLL_QUEUE, countdown=0)
        await register_poll_active(
            task,
            check_at=_poll_queue_timeout_at(),
            reason=f"{source}_has_provider_task_id",
        )
        return "recover_poll_has_provider_id"

    # 未过 deadline，且没有结果 URL / 供应商任务 ID：
    # 图片同步任务会重新进入 submit_image_task；视频/其它任务会重新创建供应商任务。
    # 这里不能投 poll，因为没有 provider_task_id/seedance_task_id 可查询。
    recoverable_create_stages = {
        ChatGenerationPipelineStage.QUEUED.value,
        ChatGenerationPipelineStage.PREPARING.value,
        ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
        ChatGenerationPipelineStage.WAITING_REMOTE.value,
        ChatGenerationPipelineStage.POLLING.value,
    }
    if task.pipeline_stage in recoverable_create_stages:
        if task.pipeline_stage not in (
            ChatGenerationPipelineStage.QUEUED.value,
            ChatGenerationPipelineStage.PREPARING.value,
            ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
        ):
            task.pipeline_stage = ChatGenerationPipelineStage.QUEUED.value
            await db.commit()

        await _remove_poll_active(task.id)
        await log_task_event(
            task,
            event_type="GENERATION_RECOVERY_ENQUEUE",
            message=f"{source} 发现任务未超时且缺少 remote_result_url/供应商任务ID，恢复投递创建队列",
            detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
        )
        chatapi_create_generation_task.apply_async(
            args=[task.id],
            queue="gen_chatapi_create",
            countdown=0,
        )
        return "recover_create_no_remote_no_provider_before_deadline"

    # result_ready 但没有 URL 是脏状态；未过 deadline 时回创建队列重新处理，过期上面已标记超时。
    if task.pipeline_stage == ChatGenerationPipelineStage.RESULT_READY.value:
        task.pipeline_stage = ChatGenerationPipelineStage.QUEUED.value
        await db.commit()
        await _remove_poll_active(task.id)
        await log_task_event(
            task,
            event_type="GENERATION_RECOVERY_ENQUEUE",
            message=f"{source} 发现 result_ready 但缺少 remote_result_url，未超时，恢复投递创建队列",
            detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
        )
        chatapi_create_generation_task.apply_async(
            args=[task.id],
            queue="gen_chatapi_create",
            countdown=0,
        )
        return "recover_create_result_ready_no_url_before_deadline"

    return f"skip_stage_{task.pipeline_stage}"


async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
    """启动时生成链路容灾扫描。

    不新增 Celery beat，不新增 worker 命令；worker 启动时由 Redis 锁保证只投递一次。
    恢复顺序：
    1. Redis poll active_index 到期任务；
    2. DB fallback 扫描 queued/creating/waiting_remote/polling/result_ready；
    3. 下载阶段仍由 recover_download_tasks_once 兜底。
    """
    checked_ids: set[str] = set()
    results: dict[str, int] = {}

    due_poll_ids = await redis_get_due_registry_ids(
        zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
        limit=int(settings.POLL_RECOVERY_BATCH_SIZE or settings.GENERATION_RECOVERY_BATCH_SIZE or 100),
        log_context="poll_active",
    )
    poll_payloads = await redis_get_registry_payloads(
        hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
        item_ids=due_poll_ids,
        log_context="poll_active",
    )

    for task_id in due_poll_ids:
        result = await db.execute(
            select(ChatGenerationTask)
            .where(
                ChatGenerationTask.id == task_id,
                ChatGenerationTask.deleted_at.is_(None),
            )
            .with_for_update()
            .limit(1)
        )
        task = result.scalar_one_or_none()
        if task is None:
            await _remove_poll_active(task_id)
            action = "clean_missing_poll_task"
        else:
            checked_ids.add(task.id)
            action = await recover_one_generation_task(
                db,
                task,
                payload=poll_payloads.get(task_id),
                source="startup_poll_redis",
            )
        results[action] = results.get(action, 0) + 1

    batch_size = int(settings.GENERATION_RECOVERY_BATCH_SIZE or settings.DOWNLOAD_RECOVERY_BATCH_SIZE or 100)
    max_rounds = max(1, int(settings.GENERATION_RECOVERY_MAX_ROUNDS or 1))
    total_db_checked = 0

    for _round in range(max_rounds):
        query_result = await db.execute(
            select(ChatGenerationTask)
            .where(
                ChatGenerationTask.deleted_at.is_(None),
                ChatGenerationTask.generation_mode.in_(list(ALLOWED_GENERATION_MODES)),
                ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
                ChatGenerationTask.pipeline_stage.in_(
                    [
                        "queued",
                        "preparing",
                        "creating_provider_task",
                        "waiting_remote",
                        "polling",
                        "result_ready",
                    ]
                ),
            )
            .order_by(ChatGenerationTask.updated_at.asc())
            .limit(batch_size)
            .with_for_update(skip_locked=True)
        )
        tasks = query_result.scalars().all()
        if not tasks:
            break

        progressed_this_round = 0
        for task in tasks:
            if task.id in checked_ids:
                continue
            action = await recover_one_generation_task(
                db,
                task,
                payload=None,
                source="startup_db",
            )
            results[action] = results.get(action, 0) + 1
            checked_ids.add(task.id)
            total_db_checked += 1
            progressed_this_round += 1

        if len(tasks) < batch_size or progressed_this_round <= 0:
            break

    return {
        "checked": len(checked_ids),
        "db_checked": total_db_checked,
        "results": results,
    }