from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

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

from app.config import settings
from app.enums.generation_status import GenerationRecordPipelineStage, GenerationStatus
from app.enums.generation_task import GenerationOwnerType
from app.models.generation_record import GenerationRecord
from app.services.generation.pipeline.owner_service import GenerationOwnerRef
from app.services.redis_registry_service import ensure_aware_utc


@dataclass(frozen=True, slots=True)
class GenerationRecordRecoveryCursor:
    owner_id: str


@dataclass(slots=True)
class GenerationRecordRecoveryBatch:
    create: list[GenerationOwnerRef]
    poll: list[GenerationOwnerRef]
    download: list[GenerationOwnerRef]
    next_cursor: GenerationRecordRecoveryCursor | None


async def find_generation_record_recovery_batch(
    db: AsyncSession,
    *,
    limit: int,
    cursor: GenerationRecordRecoveryCursor | None = None,
) -> GenerationRecordRecoveryBatch:
    """按稳定游标读取恢复分流所需列，避免大字段加载和 offset 扫描。"""
    stages = {
        GenerationRecordPipelineStage.QUEUED.value,
        GenerationRecordPipelineStage.PREPARING.value,
        GenerationRecordPipelineStage.CREATING_PROVIDER_TASK.value,
        GenerationRecordPipelineStage.WAITING_REMOTE.value,
        GenerationRecordPipelineStage.POLLING.value,
        GenerationRecordPipelineStage.RESULT_READY.value,
        GenerationRecordPipelineStage.DOWNLOAD_QUEUED.value,
        GenerationRecordPipelineStage.DOWNLOADING.value,
        GenerationRecordPipelineStage.RETRY_WAITING.value,
    }
    page_size = max(1, int(limit))
    now = datetime.now(timezone.utc)
    queue_timeout = timedelta(
        seconds=max(1, int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))
    )
    query = select(
        GenerationRecord.id,
        GenerationRecord.generation_attempt_no,
        GenerationRecord.seedance_task_id,
        GenerationRecord.remote_result_url,
        GenerationRecord.pipeline_stage,
        GenerationRecord.provider_create_lease_until,
        GenerationRecord.next_poll_at,
        GenerationRecord.poll_lease_until,
        GenerationRecord.download_enqueued_at,
        GenerationRecord.download_lease_until,
        GenerationRecord.download_next_retry_at,
    ).where(
        GenerationRecord.deleted_at.is_(None),
        GenerationRecord.status == GenerationStatus.generating.value,
        GenerationRecord.pipeline_stage.in_(stages),
    )
    if cursor is not None:
        query = query.where(GenerationRecord.id > cursor.owner_id)
    result = await db.execute(
        query.order_by(GenerationRecord.id.asc()).limit(page_size)
    )
    rows = list(result.all())

    create: list[GenerationOwnerRef] = []
    poll: list[GenerationOwnerRef] = []
    download: list[GenerationOwnerRef] = []
    for (
        owner_id,
        attempt_no,
        provider_task_id,
        remote_result_url,
        pipeline_stage,
        provider_create_lease_until,
        next_poll_at,
        poll_lease_until,
        download_enqueued_at,
        download_lease_until,
        download_next_retry_at,
    ) in rows:
        ref = GenerationOwnerRef(
            GenerationOwnerType.GENERATION_RECORD.value,
            str(owner_id),
            int(attempt_no or 1),
        )
        stage = str(pipeline_stage or "")
        if str(remote_result_url or "").strip():
            if stage == GenerationRecordPipelineStage.RESULT_READY.value:
                download.append(ref)
            elif stage == GenerationRecordPipelineStage.DOWNLOAD_QUEUED.value:
                checked_enqueued_at = ensure_aware_utc(download_enqueued_at)
                if (
                    checked_enqueued_at is None
                    or checked_enqueued_at + queue_timeout <= now
                ):
                    download.append(ref)
            elif stage == GenerationRecordPipelineStage.DOWNLOADING.value:
                if (
                    ensure_aware_utc(download_lease_until) is None
                    or ensure_aware_utc(download_lease_until) <= now
                ):
                    download.append(ref)
            elif stage == GenerationRecordPipelineStage.RETRY_WAITING.value:
                if (
                    ensure_aware_utc(download_next_retry_at) is None
                    or ensure_aware_utc(download_next_retry_at) <= now
                ):
                    download.append(ref)
        elif str(provider_task_id or "").strip():
            checked_next_poll_at = ensure_aware_utc(next_poll_at)
            checked_poll_lease_until = ensure_aware_utc(poll_lease_until)
            if (
                (checked_next_poll_at is None or checked_next_poll_at <= now)
                and (
                    checked_poll_lease_until is None
                    or checked_poll_lease_until <= now
                )
            ):
                poll.append(ref)
        else:
            if (
                ensure_aware_utc(provider_create_lease_until) is None
                or ensure_aware_utc(provider_create_lease_until) <= now
            ):
                create.append(ref)

    next_cursor = None
    if len(rows) == page_size:
        last = rows[-1]
        next_cursor = GenerationRecordRecoveryCursor(owner_id=str(last.id))
    return GenerationRecordRecoveryBatch(
        create=create,
        poll=poll,
        download=download,
        next_cursor=next_cursor,
    )