from __future__ import annotations

from collections.abc import Iterable
from datetime import datetime
from typing import Any, TypedDict

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

from app.enums.generation_task import GenerationType
from app.enums.recent_generation import (
    RECENT_GENERATION_ALL_MODULES,
    RECENT_GENERATION_CHAT_TASK_MODULES,
    RECENT_GENERATION_COMPLETED_STATUS,
    RECENT_GENERATION_MODULE_TO_TASK_MODE,
    RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE,
    RecentGenerationModuleEnum,
    RecentGenerationResourceTypeEnum,
)
from app.models.chat_generation_task import ChatGenerationTask
from app.models.generation_record import GenerationRecord
from app.models.module_generation_step import ModuleGenerationStep
from app.models.shot_replicate_segment import ShotReplicateSegment
from app.schemas.recent_generation import RecentGenerationGroupOut, RecentGenerationItemOut
from app.services.resource_signed_url_service import build_resource_signed_url

DEFAULT_RECENT_GENERATION_LIMIT = 5
MAX_RECENT_GENERATION_LIMIT = 100


class _StepLinkInfo(TypedDict):
    module_project_id: str | None
    module_step_id: str | None
    module: str | None


class _ShotLinkInfo(TypedDict):
    shot_task_set_id: str | None
    shot_segment_id: str | None


def _normalize_limit(limit: int | None) -> int:
    if limit is None:
        return DEFAULT_RECENT_GENERATION_LIMIT
    return min(max(int(limit), 1), MAX_RECENT_GENERATION_LIMIT)


def _normalize_modules(
    modules: Iterable[RecentGenerationModuleEnum] | None,
) -> list[RecentGenerationModuleEnum]:
    if not modules:
        return list(RECENT_GENERATION_ALL_MODULES)

    normalized: list[RecentGenerationModuleEnum] = []
    seen: set[RecentGenerationModuleEnum] = set()
    for module in modules:
        module_enum = RecentGenerationModuleEnum(module)
        if module_enum not in seen:
            normalized.append(module_enum)
            seen.add(module_enum)
    return normalized


def _has_url(column) -> Any:
    return and_(column.is_not(None), column != "")


def _detect_resource_type(
    gen_type: str | None,
    image_url: str | None,
    video_url: str | None,
) -> RecentGenerationResourceTypeEnum:
    gen_type_value = (gen_type or "").strip().lower()

    if gen_type_value == GenerationType.IMAGE.value and image_url:
        return RecentGenerationResourceTypeEnum.IMAGE
    if gen_type_value == GenerationType.VIDEO.value and video_url:
        return RecentGenerationResourceTypeEnum.VIDEO
    if video_url:
        return RecentGenerationResourceTypeEnum.VIDEO
    return RecentGenerationResourceTypeEnum.IMAGE


def _sign_url(url: str | None) -> str | None:
    if not url:
        return None
    return build_resource_signed_url(url)


def _build_item(
    *,
    generation_id: str,
    module: RecentGenerationModuleEnum,
    gen_type: str | None,
    image_url: str | None,
    video_url: str | None,
    video_cover_url: str | None,
    generated_time: datetime | None,
    step_info: _StepLinkInfo | None = None,
    shot_info: _ShotLinkInfo | None = None,
) -> RecentGenerationItemOut:
    resource_type = _detect_resource_type(
        gen_type=gen_type,
        image_url=image_url,
        video_url=video_url,
    )

    raw_result_url = video_url if resource_type == RecentGenerationResourceTypeEnum.VIDEO else image_url
    raw_cover_url = video_cover_url if resource_type == RecentGenerationResourceTypeEnum.VIDEO else None

    return RecentGenerationItemOut(
        generated_time=generated_time,
        result_url=_sign_url(raw_result_url),
        cover_url=_sign_url(raw_cover_url),
        module=module,
        shot_task_set_id=shot_info["shot_task_set_id"] if shot_info else None,
        shot_segment_id=shot_info["shot_segment_id"] if shot_info else None,
        module_project_id=step_info["module_project_id"] if step_info else None,
        module_step_id=step_info["module_step_id"] if step_info else None,
        generation_id=generation_id,
        resource_type=resource_type,
    )


async def _list_project_recent_items(
    db: AsyncSession,
    *,
    user_id: str,
    limit: int,
) -> list[RecentGenerationItemOut]:
    generated_time_expr = func.coalesce(
        GenerationRecord.generated_at,
        GenerationRecord.updated_at,
        GenerationRecord.created_at,
    ).label("generated_time")

    stmt = (
        select(
            GenerationRecord.id.label("generation_id"),
            GenerationRecord.gen_type.label("gen_type"),
            GenerationRecord.image_url.label("image_url"),
            GenerationRecord.video_url.label("video_url"),
            GenerationRecord.video_cover_url.label("video_cover_url"),
            generated_time_expr,
        )
        .where(
            GenerationRecord.user_id == user_id,
            GenerationRecord.deleted_at.is_(None),
            GenerationRecord.status == RECENT_GENERATION_COMPLETED_STATUS,
            or_(_has_url(GenerationRecord.image_url), _has_url(GenerationRecord.video_url)),
        )
        .order_by(generated_time_expr.desc(), GenerationRecord.created_at.desc())
        .limit(limit)
    )

    rows = (await db.execute(stmt)).mappings().all()
    return [
        _build_item(
            generation_id=row["generation_id"],
            module=RecentGenerationModuleEnum.PROJECT,
            gen_type=row["gen_type"],
            image_url=row["image_url"],
            video_url=row["video_url"],
            video_cover_url=row["video_cover_url"],
            generated_time=row["generated_time"],
        )
        for row in rows
    ]


async def _list_chat_task_recent_rows(
    db: AsyncSession,
    *,
    user_id: str,
    modules: list[RecentGenerationModuleEnum],
    limit: int,
) -> list[dict[str, Any]]:
    task_mode_values = [
        RECENT_GENERATION_MODULE_TO_TASK_MODE[module].value
        for module in modules
        if module in RECENT_GENERATION_CHAT_TASK_MODULES
    ]
    if not task_mode_values:
        return []

    generated_time_expr = func.coalesce(
        ChatGenerationTask.generated_at,
        ChatGenerationTask.updated_at,
        ChatGenerationTask.created_at,
    )

    ranked_subquery = (
        select(
            ChatGenerationTask.id.label("generation_id"),
            ChatGenerationTask.generation_mode.label("generation_mode"),
            ChatGenerationTask.gen_type.label("gen_type"),
            ChatGenerationTask.image_url.label("image_url"),
            ChatGenerationTask.video_url.label("video_url"),
            ChatGenerationTask.video_cover_url.label("video_cover_url"),
            generated_time_expr.label("generated_time"),
            func.row_number()
            .over(
                partition_by=ChatGenerationTask.generation_mode,
                order_by=(generated_time_expr.desc(), ChatGenerationTask.created_at.desc()),
            )
            .label("row_num"),
        )
        .where(
            ChatGenerationTask.user_id == user_id,
            ChatGenerationTask.deleted_at.is_(None),
            ChatGenerationTask.status == RECENT_GENERATION_COMPLETED_STATUS,
            ChatGenerationTask.generation_mode.in_(task_mode_values),
            or_(_has_url(ChatGenerationTask.image_url), _has_url(ChatGenerationTask.video_url)),
        )
        .subquery()
    )

    stmt = (
        select(ranked_subquery)
        .where(ranked_subquery.c.row_num <= limit)
        .order_by(ranked_subquery.c.generation_mode.asc(), ranked_subquery.c.generated_time.desc())
    )

    return [dict(row) for row in (await db.execute(stmt)).mappings().all()]


async def _load_step_link_map(
    db: AsyncSession,
    *,
    chat_task_ids: list[str],
) -> dict[str, _StepLinkInfo]:
    if not chat_task_ids:
        return {}

    stmt = (
        select(
            ModuleGenerationStep.chat_task_id.label("chat_task_id"),
            ModuleGenerationStep.id.label("module_step_id"),
            ModuleGenerationStep.project_id.label("module_project_id"),
            ModuleGenerationStep.module.label("module"),
            ModuleGenerationStep.is_current.label("is_current"),
            ModuleGenerationStep.updated_at.label("updated_at"),
        )
        .where(
            ModuleGenerationStep.deleted_at.is_(None),
            ModuleGenerationStep.chat_task_id.in_(chat_task_ids),
            ModuleGenerationStep.module.in_(
                [
                    RecentGenerationModuleEnum.HOT_OPENING_REPLICATE.value,
                    RecentGenerationModuleEnum.SHOT_REPLICATE.value,
                ]
            ),
        )
        .order_by(
            ModuleGenerationStep.chat_task_id.asc(),
            ModuleGenerationStep.is_current.desc(),
            ModuleGenerationStep.updated_at.desc(),
        )
    )

    link_map: dict[str, _StepLinkInfo] = {}
    rows = (await db.execute(stmt)).mappings().all()
    for row in rows:
        chat_task_id = row["chat_task_id"]
        if not chat_task_id or chat_task_id in link_map:
            continue
        link_map[chat_task_id] = {
            "module_project_id": row["module_project_id"],
            "module_step_id": row["module_step_id"],
            "module": row["module"],
        }
    return link_map


async def _load_shot_link_map(
    db: AsyncSession,
    *,
    module_project_ids: list[str],
) -> dict[str, _ShotLinkInfo]:
    if not module_project_ids:
        return {}

    stmt = (
        select(
            ShotReplicateSegment.module_project_id.label("module_project_id"),
            ShotReplicateSegment.id.label("shot_segment_id"),
            ShotReplicateSegment.task_set_id.label("shot_task_set_id"),
        )
        .where(
            ShotReplicateSegment.deleted_at.is_(None),
            ShotReplicateSegment.module_project_id.in_(module_project_ids),
        )
        .order_by(ShotReplicateSegment.updated_at.desc())
    )

    link_map: dict[str, _ShotLinkInfo] = {}
    rows = (await db.execute(stmt)).mappings().all()
    for row in rows:
        module_project_id = row["module_project_id"]
        if not module_project_id or module_project_id in link_map:
            continue
        link_map[module_project_id] = {
            "shot_task_set_id": row["shot_task_set_id"],
            "shot_segment_id": row["shot_segment_id"],
        }
    return link_map


async def _build_chat_task_group_items(
    db: AsyncSession,
    *,
    rows: list[dict[str, Any]],
) -> dict[RecentGenerationModuleEnum, list[RecentGenerationItemOut]]:
    grouped: dict[RecentGenerationModuleEnum, list[RecentGenerationItemOut]] = {
        RecentGenerationModuleEnum.CHAT_AI: [],
        RecentGenerationModuleEnum.HOT_OPENING_REPLICATE: [],
        RecentGenerationModuleEnum.SHOT_REPLICATE: [],
    }
    if not rows:
        return grouped

    module_task_rows: list[dict[str, Any]] = []
    module_chat_task_ids: list[str] = []
    for row in rows:
        module = RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE.get(row["generation_mode"])
        if module in (
            RecentGenerationModuleEnum.HOT_OPENING_REPLICATE,
            RecentGenerationModuleEnum.SHOT_REPLICATE,
        ):
            module_task_rows.append(row)
            module_chat_task_ids.append(row["generation_id"])

    step_link_map = await _load_step_link_map(db, chat_task_ids=module_chat_task_ids)

    shot_project_ids = [
        step_info["module_project_id"]
        for row in module_task_rows
        if (step_info := step_link_map.get(row["generation_id"]))
        and step_info.get("module") == RecentGenerationModuleEnum.SHOT_REPLICATE.value
        and step_info.get("module_project_id")
    ]
    shot_link_map = await _load_shot_link_map(db, module_project_ids=shot_project_ids)

    for row in rows:
        module = RECENT_GENERATION_TASK_MODE_VALUE_TO_MODULE.get(row["generation_mode"])
        if not module:
            continue

        step_info = step_link_map.get(row["generation_id"])
        shot_info = None
        if module == RecentGenerationModuleEnum.SHOT_REPLICATE and step_info:
            module_project_id = step_info.get("module_project_id")
            if module_project_id:
                shot_info = shot_link_map.get(module_project_id)

        grouped[module].append(
            _build_item(
                generation_id=row["generation_id"],
                module=module,
                gen_type=row["gen_type"],
                image_url=row["image_url"],
                video_url=row["video_url"],
                video_cover_url=row["video_cover_url"],
                generated_time=row["generated_time"],
                step_info=step_info,
                shot_info=shot_info,
            )
        )

    return grouped


async def list_recent_generations(
    db: AsyncSession,
    *,
    user_id: str,
    modules: Iterable[RecentGenerationModuleEnum] | None = None,
    limit: int | None = None,
) -> RecentGenerationGroupOut:
    """获取当前用户各模块最近生成成功的图片/视频记录。"""

    normalized_limit = _normalize_limit(limit)
    normalized_modules = _normalize_modules(modules)

    response = RecentGenerationGroupOut()

    if RecentGenerationModuleEnum.PROJECT in normalized_modules:
        response.project = await _list_project_recent_items(
            db,
            user_id=user_id,
            limit=normalized_limit,
        )

    chat_modules = [module for module in normalized_modules if module in RECENT_GENERATION_CHAT_TASK_MODULES]
    if chat_modules:
        chat_rows = await _list_chat_task_recent_rows(
            db,
            user_id=user_id,
            modules=chat_modules,
            limit=normalized_limit,
        )
        chat_grouped = await _build_chat_task_group_items(db, rows=chat_rows)
        response.chat_ai = chat_grouped[RecentGenerationModuleEnum.CHAT_AI]
        response.hot_opening_replicate = chat_grouped[RecentGenerationModuleEnum.HOT_OPENING_REPLICATE]
        response.shot_replicate = chat_grouped[RecentGenerationModuleEnum.SHOT_REPLICATE]

    return response