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.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_provider_call, log_task_event
from app.services.generation_module_hook_service import notify_chat_generation_task_finished
from app.services.generation_provider_service import poll_provider_task
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")

ALLOWED_GENERATION_MODES = {"chatapi_async", "hot_opening_replicate", "shot_replicate"}
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 ("completed", "failed") or task.pipeline_stage in (
        "done",
        "failed",
        "timeout",
        "download_failed",
    )


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 != "generating":
        await remove_download_active(task.id)
        return "clean_not_generating"
    if not task.remote_result_url:
        return "skip_no_remote_result_url"

    stage = task.pipeline_stage
    redis_payload = payload or {}

    if stage == "result_ready":
        await log_task_event(
            task,
            event_type="DOWNLOAD_RECOVERY_ENQUEUE",
            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="DOWNLOAD_RECOVERY_ENQUEUE",
                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="DOWNLOAD_RECOVERY_ENQUEUE",
                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="DOWNLOAD_RECOVERY_ENQUEUE",
                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_(["chatapi_async", "hot_opening_replicate", "shot_replicate"]),
            ChatGenerationTask.status == "generating",
            ChatGenerationTask.remote_result_url.is_not(None),
            ChatGenerationTask.pipeline_stage.in_(
                ["result_ready", "download_queued", "downloading", "retry_waiting"]
            ),
        )
        .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="timeout",
    )
    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="timeout",
    )
    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="failed",
    )
    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 _try_final_poll_before_timeout(db: AsyncSession, task: ChatGenerationTask) -> str:
    """超时前最后查一次供应商，避免 Celery 中断导致本地假超时。

    如果供应商已经成功，继续进入下载；如果仍 running 或查询失败，再按超时处理。
    """
    from app.tasks.generation_download_tasks import enqueue_download_task

    if not (task.provider_task_id or task.seedance_task_id):
        return await _mark_timeout(db, task)

    try:
        poll_result = await poll_provider_task(db, task)
        status = poll_result.get("status")
        response_data = poll_result.get("response_data")
    except Exception as exc:
        await log_task_event(
            task,
            event_type="FINAL_POLL_BEFORE_TIMEOUT_ERROR",
            message=str(exc),
        )
        return await _mark_timeout(db, task)

    try:
        provider_response = json.loads(response_data or "{}")
    except Exception:
        provider_response = {"raw": response_data}

    snapshot = _engine_snapshot(task)
    await log_provider_call(
        task,
        provider=snapshot.get("provider") or "ark",
        api_type=f"{task.gen_type}_final_poll_before_timeout",
        model=snapshot.get("model_name"),
        engine_id=task.engine_id,
        status="success",
        provider_task_id=task.seedance_task_id or task.provider_task_id,
        response_data=provider_response,
    )

    if _is_success(status):
        if task.gen_type == "image":
            task.remote_result_url = poll_result.get("image_url")
            task.image_tokens_used = poll_result.get("image_tokens", 0) or 0
        else:
            task.remote_result_url = poll_result.get("video_url")
            task.video_tokens_used = poll_result.get("video_tokens", 0) or 0

        task.provider_response_json = response_data
        if not task.remote_result_url:
            return await _mark_failed(
                db,
                task,
                error_message="供应商任务成功但未返回结果URL",
                detail=poll_result,
            )

        task.pipeline_stage = "result_ready"
        task.retry_count = 0
        await db.commit()
        await _remove_poll_active(task.id)
        await log_task_event(
            task,
            event_type="POLL_SUCCESS_AFTER_TIMEOUT_RECOVERY",
            to_stage="result_ready",
            detail=poll_result,
        )
        await enqueue_download_task(db, task, recover=True, reason="final_poll_before_timeout_success")
        return "recover_timeout_success_to_download"

    if _is_failed(status):
        task.provider_response_json = response_data
        return await _mark_failed(
            db,
            task,
            error_message=poll_result.get("error") or f"供应商任务失败: {status}",
            detail=poll_result,
        )

    await log_task_event(
        task,
        event_type="FINAL_POLL_BEFORE_TIMEOUT_PENDING",
        message=f"status={status}",
        detail=poll_result,
    )
    return await _mark_timeout(db, task)


async def recover_one_generation_task(
    db: AsyncSession,
    task: ChatGenerationTask,
    *,
    payload: dict[str, Any] | None = None,
    source: str = "startup_db",
) -> str:
    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 != "generating":
        await _remove_poll_active(task.id)
        return "clean_not_generating"

    if task.deadline_at and _is_expired(task.deadline_at, current_time):
        if task.pipeline_stage in ("waiting_remote", "polling"):
            return await _try_final_poll_before_timeout(db, task)
        return await _mark_timeout(db, task)

    if task.pipeline_stage in ("queued", "preparing", "creating_provider_task"):
        if task.provider_task_id or task.seedance_task_id:
            task.pipeline_stage = "waiting_remote"
            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}_create_stage_has_provider_id",
            )
            return "recover_poll_from_create_stage"

        await log_task_event(
            task,
            event_type="GENERATION_RECOVERY_ENQUEUE",
            message=f"{source} 发现创建阶段任务未完成，恢复投递创建队列",
            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"

    if task.pipeline_stage in ("waiting_remote", "polling"):
        if task.remote_result_url:
            await _remove_poll_active(task.id)
            await enqueue_download_task(
                db,
                task,
                recover=True,
                reason=f"{source}_waiting_remote_has_result",
            )
            return "recover_waiting_has_result"

        if task.provider_task_id or task.seedance_task_id:
            await log_task_event(
                task,
                event_type="GENERATION_RECOVERY_ENQUEUE",
                message=f"{source} 发现远程等待/轮询阶段任务未完成，恢复投递轮询队列",
                detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
            )
            task.pipeline_stage = "waiting_remote"
            await db.commit()
            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}_recover_poll",
            )
            return "recover_poll"

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

    if task.pipeline_stage == "result_ready":
        await _remove_poll_active(task.id)
        if task.remote_result_url:
            await enqueue_download_task(
                db,
                task,
                recover=True,
                reason=f"{source}_generation_result_ready",
            )
            return "recover_result_ready"
        return "skip_result_ready_no_url"

    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_(["chatapi_async", "hot_opening_replicate", "shot_replicate"]),
                ChatGenerationTask.status == "generating",
                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

    # 下载阶段单独跑 DB fallback。
    download_result = await recover_download_tasks_once(db)
    return {
        "checked": len(checked_ids),
        "db_checked": total_db_checked,
        "results": results,
        "download_recovery": download_result,
    }
