from __future__ import annotations

import asyncio

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

from app.core.db import async_session_maker
from app.core.task_manager import DeliveryMode, SqlAlchemyTaskStore, TaskManager
from app.core.task_manager.types import TaskStatus
from app.core.tasks import ImageGenerationInput, ImageGenerationResult, ImageGenerationTask, ProviderConfig
from app.models.studio import (
    ActorImage,
    AssetQualityLevel,
    AssetViewAngle,
    CharacterImage,
    CostumeImage,
    PropImage,
    SceneImage,
    ShotDetail,
    ShotFrameImage,
)
from app.models.task_links import GenerationTaskLink
from app.models.types import FileUsageKind
from app.schemas.common import ApiResponse, success_response
from app.services.studio.file_usages import (
    first_project_id_for_actor,
    first_project_id_for_costume,
    first_project_id_for_prop,
    first_project_id_for_scene,
    sync_usage_from_character,
    sync_usage_from_shot_context,
    upsert_file_usage,
)
from app.services.studio.shot_status import mark_shot_generating, recompute_shot_status
from app.services.studio.image_tasks import load_provider_config, resolve_image_model
from app.utils.files import create_file_from_url_or_b64
from app.api.v1.routes.film.common import TaskCreated, _CreateOnlyTask


async def _persist_images_to_assets(
    session: AsyncSession,
    *,
    task_id: str,
    relation_type: str,
    relation_entity_id: str,
    result: ImageGenerationResult,
) -> None:
    """将图片生成结果落库到 FileItem 与业务图片表。"""
    images = result.images or []
    if not images:
        return

    item = images[0]
    if not item.url:
        return

    file_obj = await create_file_from_url_or_b64(
        session,
        url=item.url,
        name=f"{relation_type}-{relation_entity_id}",
        prefix=f"generated-images/{relation_type}/{relation_entity_id}",
    )
    file_id = file_obj.id

    link_stmt = (
        select(GenerationTaskLink)
        .where(
            GenerationTaskLink.task_id == task_id,
            GenerationTaskLink.relation_type == relation_type,
            GenerationTaskLink.relation_entity_id == relation_entity_id,
        )
        .limit(1)
    )
    link_row = (await session.execute(link_stmt)).scalars().first()
    if link_row is not None:
        link_row.file_id = file_id

    if relation_type == "actor_image":
        image_row = await session.get(ActorImage, int(relation_entity_id))
        if image_row is not None:
            image_row.file_id = file_id
            pid = await first_project_id_for_actor(session, image_row.actor_id)
            if pid:
                await upsert_file_usage(
                    session,
                    file_id=file_id,
                    project_id=pid,
                    chapter_id=None,
                    shot_id=None,
                    usage_kind=FileUsageKind.asset_image,
                    source_ref=f"actor_image:{image_row.id}",
                )
    elif relation_type == "scene_image":
        image_row = await session.get(SceneImage, int(relation_entity_id))
        if image_row is not None:
            image_row.file_id = file_id
            pid = await first_project_id_for_scene(session, image_row.scene_id)
            if pid:
                await upsert_file_usage(
                    session,
                    file_id=file_id,
                    project_id=pid,
                    chapter_id=None,
                    shot_id=None,
                    usage_kind=FileUsageKind.asset_image,
                    source_ref=f"scene_image:{image_row.id}",
                )
    elif relation_type == "prop_image":
        image_row = await session.get(PropImage, int(relation_entity_id))
        if image_row is not None:
            image_row.file_id = file_id
            pid = await first_project_id_for_prop(session, image_row.prop_id)
            if pid:
                await upsert_file_usage(
                    session,
                    file_id=file_id,
                    project_id=pid,
                    chapter_id=None,
                    shot_id=None,
                    usage_kind=FileUsageKind.asset_image,
                    source_ref=f"prop_image:{image_row.id}",
                )
    elif relation_type == "costume_image":
        image_row = await session.get(CostumeImage, int(relation_entity_id))
        if image_row is not None:
            image_row.file_id = file_id
            pid = await first_project_id_for_costume(session, image_row.costume_id)
            if pid:
                await upsert_file_usage(
                    session,
                    file_id=file_id,
                    project_id=pid,
                    chapter_id=None,
                    shot_id=None,
                    usage_kind=FileUsageKind.asset_image,
                    source_ref=f"costume_image:{image_row.id}",
                )
    elif relation_type == "character_image":
        image_row = await session.get(CharacterImage, int(relation_entity_id))
        if image_row is not None:
            image_row.file_id = file_id
            await sync_usage_from_character(
                session,
                file_id=file_id,
                character_id=image_row.character_id,
                usage_kind=FileUsageKind.character_image,
                source_ref=f"character_image:{image_row.id}",
            )
    elif relation_type == "character":
        character_id = relation_entity_id
        stmt_ci = (
            select(CharacterImage)
            .where(
                CharacterImage.character_id == character_id,
                CharacterImage.quality_level == AssetQualityLevel.low,
                CharacterImage.view_angle == AssetViewAngle.front,
            )
            .order_by(CharacterImage.id.asc())
            .limit(1)
        )
        ci = (await session.execute(stmt_ci)).scalars().first()
        if ci is not None:
            ci.file_id = file_id
            ci.format = getattr(ci, "format", "") or "png"
        else:
            ci = CharacterImage(
                character_id=character_id,
                file_id=file_id,
                quality_level=AssetQualityLevel.low,
                view_angle=AssetViewAngle.front,
                width=None,
                height=None,
                format="png",
                is_primary=True,
            )
            session.add(ci)

        if ci is not None and getattr(ci, "is_primary", False) is True and getattr(ci, "id", None) is not None:
            stmt_clear = (
                CharacterImage.__table__.update()  # type: ignore[attr-defined]
                .where(CharacterImage.character_id == character_id, CharacterImage.id != ci.id)
                .values(is_primary=False)
            )
            await session.execute(stmt_clear)
        await session.flush()
        if ci is not None:
            await sync_usage_from_character(
                session,
                file_id=file_id,
                character_id=character_id,
                usage_kind=FileUsageKind.character_image,
                source_ref=f"character_image:{ci.id}",
            )
    elif relation_type == "shot_frame_image":
        image_row = await session.get(ShotFrameImage, int(relation_entity_id))
        if image_row is not None:
            image_row.file_id = file_id
            detail = await session.get(ShotDetail, image_row.shot_detail_id)
            if detail is not None:
                await sync_usage_from_shot_context(
                    session,
                    file_id=file_id,
                    shot_id=detail.id,
                    usage_kind=FileUsageKind.shot_frame,
                    source_ref=f"shot_frame_image:{image_row.id}",
                )


async def _resolve_related_shot_id(
    session: AsyncSession,
    *,
    relation_type: str,
    relation_entity_id: str,
) -> str | None:
    """仅解析和镜头直接相关的生成任务。"""
    if relation_type != "shot_frame_image":
        return None
    image_row = await session.get(ShotFrameImage, int(relation_entity_id))
    if image_row is None:
        return None
    return image_row.shot_detail_id


async def create_image_task_and_link(
    *,
    db: AsyncSession,
    model_id: str | None,
    relation_type: str,
    relation_entity_id: str,
    prompt: str,
    images: list[dict[str, str]] | None = None,
) -> TaskCreated:
    """创建图片生成任务，并建立任务关联。"""
    store = SqlAlchemyTaskStore(db)
    tm = TaskManager(store=store, strategies={})

    model = await resolve_image_model(db, model_id)
    provider_cfg = await load_provider_config(db, model.provider_id)

    run_args: dict = {
        "provider": provider_cfg.provider,
        "api_key": provider_cfg.api_key,
        "base_url": provider_cfg.base_url,
        "input": {
            "prompt": prompt,
            "model": model.name,
        },
    }
    if images:
        run_args["input"]["images"] = images

    task_record = await tm.create(
        task=_CreateOnlyTask(),
        mode=DeliveryMode.async_polling,
        run_args=run_args,
    )

    db.add(
        GenerationTaskLink(
            task_id=task_record.id,
            resource_type="image",
            relation_type=relation_type,
            relation_entity_id=relation_entity_id,
        )
    )
    related_shot_id = await _resolve_related_shot_id(
        db,
        relation_type=relation_type,
        relation_entity_id=relation_entity_id,
    )
    if related_shot_id:
        await mark_shot_generating(db, shot_id=related_shot_id)
    await db.commit()

    async def _runner(task_id: str, args: dict) -> None:
        async with async_session_maker() as session:
            try:
                store = SqlAlchemyTaskStore(session)
                await store.set_status(task_id, TaskStatus.running)
                await store.set_progress(task_id, 10)

                provider = str(args.get("provider") or "")
                api_key = str(args.get("api_key") or "")
                base_url = args.get("base_url")
                input_dict = dict(args.get("input") or {})

                task = ImageGenerationTask(
                    provider_config=ProviderConfig(
                        provider=provider,  # type: ignore[arg-type]
                        api_key=api_key,
                        base_url=base_url,
                    ),
                    input_=ImageGenerationInput.model_validate(input_dict),
                )
                await task.run()
                result = await task.get_result()
                if result is None:
                    raise RuntimeError("Image generation task returned no result")

                await store.set_result(task_id, result.model_dump())
                await _persist_images_to_assets(
                    session,
                    task_id=task_id,
                    relation_type=relation_type,
                    relation_entity_id=relation_entity_id,
                    result=result,
                )
                await store.set_progress(task_id, 100)
                await store.set_status(task_id, TaskStatus.succeeded)
                related_shot_id = await _resolve_related_shot_id(
                    session,
                    relation_type=relation_type,
                    relation_entity_id=relation_entity_id,
                )
                if related_shot_id:
                    await recompute_shot_status(session, shot_id=related_shot_id)
                await session.commit()
            except Exception as exc:  # noqa: BLE001
                await session.rollback()
                async with async_session_maker() as s2:
                    store = SqlAlchemyTaskStore(s2)
                    await store.set_error(task_id, str(exc))
                    await store.set_status(task_id, TaskStatus.failed)
                    related_shot_id = await _resolve_related_shot_id(
                        s2,
                        relation_type=relation_type,
                        relation_entity_id=relation_entity_id,
                    )
                    if related_shot_id:
                        await recompute_shot_status(s2, shot_id=related_shot_id)
                    await s2.commit()

    asyncio.create_task(_runner(task_record.id, run_args))
    return TaskCreated(task_id=task_record.id)
