from __future__ import annotations

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

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

from app.config import settings
from app.enums.celery_queue import CeleryQueue
from app.enums.llm_billing import LlmBillingLedgerState
from app.enums.shot_replicate import ShotSplitStatusEnum
from app.models.shot_replicate_segment import ShotReplicateSegment
from app.models.shot_replicate_task_set import ShotReplicateTaskSet
from app.services.shot_replicate_taskset_service import (
    build_segment_analysis_billing_context,
    build_task_set_analysis_billing_context,
    refresh_task_set_split_summaries,
)
from app.services.llm_billing import get_llm_ledger_states
from app.services.celery_runtime.runtime_service import runtime_lock_values
from app.tasks.celery_app import celery_app

logger = logging.getLogger("video_gen")


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


def _ensure_aware(value: datetime | None) -> datetime | None:
    if value is None:
        return None
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


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


def _queue_timeout(segment: ShotReplicateSegment, now: datetime | None = None) -> bool:
    enqueued_at = _ensure_aware(segment.split_enqueued_at)
    if enqueued_at is None:
        return True
    return enqueued_at + timedelta(seconds=int(settings.SHOT_SPLIT_PENDING_TIMEOUT_SECONDS or 300)) <= (now or _now())


async def recover_shot_split_tasks_once(db: AsyncSession) -> dict[str, Any]:
    """批量恢复拆镜 ffmpeg 任务；只接管执行锁消失且业务租约到期的记录。"""
    from app.tasks.shot_replicate_tasks import split_one_segment

    batch_size = max(1, int(settings.SHOT_SPLIT_RECOVERY_BATCH_SIZE or 50))
    current_time = _now()
    result = await db.execute(
        select(ShotReplicateSegment)
        .where(
            ShotReplicateSegment.deleted_at.is_(None),
            ShotReplicateSegment.split_status.in_(
                [
                    ShotSplitStatusEnum.PENDING.value,
                    ShotSplitStatusEnum.PROCESSING.value,
                    ShotSplitStatusEnum.RETRY_WAITING.value,
                ]
            ),
        )
        .order_by(ShotReplicateSegment.updated_at.asc(), ShotReplicateSegment.id.asc())
        .limit(batch_size)
        .with_for_update(skip_locked=True)
    )
    segments = list(result.scalars().all())
    due_segments: list[ShotReplicateSegment] = []
    results: dict[str, int] = {}
    for segment in segments:
        if segment.split_status == ShotSplitStatusEnum.PENDING.value:
            due = _queue_timeout(segment, current_time)
        elif segment.split_status == ShotSplitStatusEnum.PROCESSING.value:
            due = _expired(segment.split_lease_until, current_time)
        else:
            due = _expired(segment.split_next_retry_at, current_time)
        if not due:
            key = f"skip_{segment.split_status}_not_due"
            results[key] = results.get(key, 0) + 1
            continue
        due_segments.append(segment)

    lock_key_by_id: dict[str, str] = {}
    for segment in due_segments:
        current_attempt = max(0, int(segment.split_retry_count or 0))
        active_attempt = (
            max(1, current_attempt)
            if segment.split_status == ShotSplitStatusEnum.PROCESSING.value
            else current_attempt + 1
        )
        lock_key_by_id[str(segment.id)] = (
            f"{settings.SHOT_SPLIT_LOCK_KEY_PREFIX}:{segment.id}:attempt:{active_attempt}"
        )
    live_locks = await runtime_lock_values(list(lock_key_by_id.values()))

    dispatches: list[tuple[str, int]] = []
    touched_task_set_ids: set[str] = set()
    for segment in due_segments:
        lock_key = lock_key_by_id[str(segment.id)]
        if live_locks.get(lock_key):
            results["skip_live_runtime_lock"] = results.get("skip_live_runtime_lock", 0) + 1
            continue
        touched_task_set_ids.add(str(segment.task_set_id))
        if int(segment.split_retry_count or 0) >= int(settings.SHOT_SPLIT_MAX_RETRY_COUNT or 3):
            segment.split_status = ShotSplitStatusEnum.FAILED.value
            segment.split_last_error = segment.split_last_error or "恢复时超过最大重试次数"
            segment.split_claim_token = None
            segment.split_lease_until = None
            segment.split_next_retry_at = None
            results["mark_failed_max_retry"] = results.get("mark_failed_max_retry", 0) + 1
            continue

        segment.split_status = ShotSplitStatusEnum.PENDING.value
        segment.split_claim_token = None
        segment.split_enqueued_at = current_time
        segment.split_lease_until = None
        segment.split_next_retry_at = None
        next_attempt = int(segment.split_retry_count or 0) + 1
        dispatches.append((str(segment.id), next_attempt))
        results["recover_db"] = results.get("recover_db", 0) + 1

    await refresh_task_set_split_summaries(db, touched_task_set_ids)
    await db.commit()

    enqueue_failed = 0
    if celery_app:
        for segment_id, next_attempt in dispatches:
            try:
                split_one_segment.apply_async(
                    args=[segment_id],
                    queue=CeleryQueue.GEN_SHOT_SPLIT.value,
                    priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
                    countdown=0,
                    task_id=f"shot-split:{segment_id}:attempt:{next_attempt}",
                )
            except Exception:
                enqueue_failed += 1
    if enqueue_failed:
        results["enqueue_failed"] = enqueue_failed
    return {"checked": len(segments), "results": results}


async def recover_shot_analysis_tasks_once(db: AsyncSession) -> dict[str, Any]:
    """恢复长视频分析；只接管执行锁消失且数据库租约已到期的记录。"""
    from app.enums.shot_replicate import (
        ShotAnalysisStatusEnum,
        ShotSegmentAnalysisStatusEnum,
        ShotTaskSetStatusEnum,
    )
    from app.services.celery_runtime.runtime_service import runtime_lock_values
    from app.tasks.shot_replicate_tasks import analyze_custom_segment_video, analyze_original_video

    now = _now()
    queue_cutoff = now - timedelta(seconds=max(60, int(settings.MODULE_ASYNC_QUEUE_TIMEOUT_SECONDS or 300)))
    batch_size = max(1, int(settings.MODULE_ASYNC_RECOVERY_BATCH_SIZE or 100))

    task_set_result = await db.execute(
        select(ShotReplicateTaskSet)
        .where(
            ShotReplicateTaskSet.deleted_at.is_(None),
            ShotReplicateTaskSet.analysis_status.in_([
                ShotAnalysisStatusEnum.PENDING.value,
                ShotAnalysisStatusEnum.PROCESSING.value,
            ]),
            or_(
                (
                    (ShotReplicateTaskSet.analysis_status == ShotAnalysisStatusEnum.PENDING.value)
                    & (ShotReplicateTaskSet.updated_at <= queue_cutoff)
                ),
                (
                    (ShotReplicateTaskSet.analysis_status == ShotAnalysisStatusEnum.PROCESSING.value)
                    & or_(
                        ShotReplicateTaskSet.analysis_lease_until <= now,
                        (
                            ShotReplicateTaskSet.analysis_lease_until.is_(None)
                            & (ShotReplicateTaskSet.updated_at <= queue_cutoff)
                        ),
                    )
                ),
            ),
        )
        .order_by(ShotReplicateTaskSet.updated_at.asc(), ShotReplicateTaskSet.id.asc())
        .limit(batch_size)
        .with_for_update(skip_locked=True)
    )
    task_sets = list(task_set_result.scalars().all())

    remaining = max(0, batch_size - len(task_sets))
    segments: list[ShotReplicateSegment] = []
    if remaining:
        segment_result = await db.execute(
            select(ShotReplicateSegment)
            .where(
                ShotReplicateSegment.deleted_at.is_(None),
                ShotReplicateSegment.segment_video_url.is_not(None),
                ShotReplicateSegment.analysis_status.in_([
                    ShotSegmentAnalysisStatusEnum.PENDING.value,
                    ShotSegmentAnalysisStatusEnum.PROCESSING.value,
                ]),
                or_(
                    (
                        (ShotReplicateSegment.analysis_status == ShotSegmentAnalysisStatusEnum.PENDING.value)
                        & (ShotReplicateSegment.updated_at <= queue_cutoff)
                    ),
                    (
                        (ShotReplicateSegment.analysis_status == ShotSegmentAnalysisStatusEnum.PROCESSING.value)
                        & or_(
                            ShotReplicateSegment.analysis_lease_until <= now,
                            (
                                ShotReplicateSegment.analysis_lease_until.is_(None)
                                & (ShotReplicateSegment.updated_at <= queue_cutoff)
                            ),
                        )
                    ),
                ),
            )
            .order_by(ShotReplicateSegment.updated_at.asc(), ShotReplicateSegment.id.asc())
            .limit(remaining)
            .with_for_update(skip_locked=True)
        )
        segments = list(segment_result.scalars().all())

    billing_contexts = [
        *(build_task_set_analysis_billing_context(item) for item in task_sets),
        *(build_segment_analysis_billing_context(item) for item in segments),
    ]
    # 场景积分消费配置变更后仍需识别并继续处理已存在的有效消费；缺少消费记录时禁止恢复。
    billing_states = await get_llm_ledger_states(db, billing_contexts)

    lock_keys: list[str] = []
    task_set_lock_keys: dict[str, str] = {}
    segment_lock_keys: dict[str, str] = {}
    for item in task_sets:
        attempt = max(1, int(item.analysis_attempt_no or 1))
        key = f"{settings.SHOT_ANALYSIS_LOCK_KEY_PREFIX}:shot_task_set:{item.id}:attempt:{attempt}"
        task_set_lock_keys[str(item.id)] = key
        lock_keys.append(key)
    for item in segments:
        attempt = max(1, int(item.analysis_attempt_no or 1))
        key = f"{settings.SHOT_ANALYSIS_LOCK_KEY_PREFIX}:shot_segment:{item.id}:attempt:{attempt}"
        segment_lock_keys[str(item.id)] = key
        lock_keys.append(key)
    live_locks = await runtime_lock_values(lock_keys)

    dispatches: list[tuple[str, str, int]] = []
    results: dict[str, int] = {}
    for item in task_sets:
        key = task_set_lock_keys[str(item.id)]
        if live_locks.get(key):
            results["skip_live_task_set_lock"] = results.get("skip_live_task_set_lock", 0) + 1
            continue
        context = build_task_set_analysis_billing_context(item)
        validation = billing_states.get(context.charge_biz_key)
        can_execute = bool(validation and validation.can_execute)
        state = validation.state.value if validation else LlmBillingLedgerState.MISSING.value
        if validation and validation.state == LlmBillingLedgerState.MISSING:
            state = LlmBillingLedgerState.INVALID.value
        if not can_execute:
            item.analysis_status = ShotAnalysisStatusEnum.FAILED.value
            item.status = ShotTaskSetStatusEnum.ANALYSIS_FAILED.value
            item.analysis_claim_token = None
            item.analysis_started_at = None
            item.analysis_lease_until = None
            item.analysis_error_message = f"LLM账务状态异常（{state}），恢复任务已终止"
            result_key = f"task_set_billing_{state}"
            results[result_key] = results.get(result_key, 0) + 1
            continue
        attempt = max(1, int(item.analysis_attempt_no or 1))
        item.analysis_status = ShotAnalysisStatusEnum.PENDING.value
        item.status = ShotTaskSetStatusEnum.PENDING_ANALYSIS.value
        item.analysis_claim_token = None
        item.analysis_started_at = None
        item.analysis_lease_until = None
        dispatches.append(("task_set", str(item.id), attempt))
        results["task_set_recovered"] = results.get("task_set_recovered", 0) + 1

    for item in segments:
        key = segment_lock_keys[str(item.id)]
        if live_locks.get(key):
            results["skip_live_segment_lock"] = results.get("skip_live_segment_lock", 0) + 1
            continue
        context = build_segment_analysis_billing_context(item)
        validation = billing_states.get(context.charge_biz_key)
        can_execute = bool(validation and validation.can_execute)
        state = validation.state.value if validation else LlmBillingLedgerState.MISSING.value
        if validation and validation.state == LlmBillingLedgerState.MISSING:
            state = LlmBillingLedgerState.INVALID.value
        if not can_execute:
            item.analysis_status = ShotSegmentAnalysisStatusEnum.FAILED.value
            item.analysis_claim_token = None
            item.analysis_started_at = None
            item.analysis_lease_until = None
            item.analysis_error_message = f"LLM账务状态异常（{state}），恢复任务已终止"
            result_key = f"segment_billing_{state}"
            results[result_key] = results.get(result_key, 0) + 1
            continue
        attempt = max(1, int(item.analysis_attempt_no or 1))
        item.analysis_status = ShotSegmentAnalysisStatusEnum.PENDING.value
        item.analysis_claim_token = None
        item.analysis_started_at = None
        item.analysis_lease_until = None
        dispatches.append(("segment", str(item.id), attempt))
        results["segment_recovered"] = results.get("segment_recovered", 0) + 1

    await db.commit()
    for owner_type, owner_id, attempt in dispatches:
        try:
            if owner_type == "task_set":
                analyze_original_video.apply_async(
                    args=[owner_id, attempt],
                    queue=CeleryQueue.GEN_SHOT_ANALYSIS.value,
                    countdown=0,
                    priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
                    task_id=f"shot-analysis:task-set:{owner_id}:attempt:{attempt}",
                )
            else:
                analyze_custom_segment_video.apply_async(
                    args=[owner_id, attempt],
                    queue=CeleryQueue.GEN_SHOT_ANALYSIS.value,
                    countdown=0,
                    priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
                    task_id=f"shot-analysis:segment:{owner_id}:attempt:{attempt}",
                )
            results["enqueue_success"] = results.get("enqueue_success", 0) + 1
        except Exception:
            logger.exception(
                "恢复投递拆镜分析失败。owner_type=%s owner_id=%s attempt=%s",
                owner_type,
                owner_id,
                attempt,
            )
            results["enqueue_failed"] = results.get("enqueue_failed", 0) + 1
    return {"checked": len(task_sets) + len(segments), "results": results}