from __future__ import annotations

from fastapi import APIRouter, Body, Depends, HTTPException, Path
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from app.dependencies import get_current_user, get_db
from app.models.chat_generation_task import ChatGenerationTask
from app.models.user import User
from app.schemas.module_generation_v2 import (
    ModuleVideoPromptRetryV2,
    ModuleVideoPromptSchemaUpdateV2,
    ShotReplicateProjectCreateV2,
)
from app.schemas.shot_replicate import ShotReplicateActionOut, ShotReplicateDeleteOut, ShotReplicateTaskDetailOut
from app.services.generation.pipeline.enqueue_service import enqueue_generation_create
from app.services.llm_billing import LlmBillingContext, log_celery_dispatch_compensated
from app.services.module_async_recovery_service import OBJECT_MODULE_STEP, has_live_object_lock
from app.services.module_generation_v2.config import SHOT_REPLICATE_V2
from app.services.module_generation_v2.dispatch_service import (
    dispatch_video_prompt_v2,
    ensure_v2_celery_enabled,
)
from app.services.module_generation_v2.flow_service import (
    build_v2_video_prompt_billing_context,
    create_shot_replicate_project_v2,
    delete_project_v2,
    generate_video_from_prompt_v2,
    get_v2_project_for_user,
    is_project_idempotency_conflict,
    mark_video_prompt_dispatch_failed_v2,
    rebuild_video_prompt_step_v2,
    update_video_prompt_schema_v2,
)
from app.services.shot_replicate_flow_service import project_to_detail_out
from app.services.shot_replicate_taskset_service import get_segment_for_user
from app.services.upload_resource import cleanup_upload_resource_files_after_commit

router = APIRouter(prefix="/shot-replications", tags=["shot-replications-v2"])


def _dispatch_context(*, user_id: str, project_id: str, step_id: str, step_version: int) -> LlmBillingContext:
    context = build_v2_video_prompt_billing_context(
        user_id=user_id,
        project_id=project_id,
        step_id=step_id,
        step_version=step_version,
        module=SHOT_REPLICATE_V2.module,
        display_name=SHOT_REPLICATE_V2.display_name,
    )
    context.celery_task_id = f"module-v2-video-prompt:{step_id}"
    return context


async def _detail(db: AsyncSession, current_user: User, project_id: str) -> ShotReplicateTaskDetailOut:
    project = await get_v2_project_for_user(
        db,
        config=SHOT_REPLICATE_V2,
        project_id=project_id,
        current_user=current_user,
    )
    return await project_to_detail_out(db, project)


async def _dispatch_or_mark_failed(
    db: AsyncSession,
    *,
    project_id: str,
    step_id: str,
    billing_context: LlmBillingContext,
) -> None:
    dispatch = await dispatch_video_prompt_v2(
        config=SHOT_REPLICATE_V2,
        project_id=project_id,
        step_id=step_id,
        billing_context=billing_context,
    )
    if dispatch.recoverable:
        return
    if await has_live_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id):
        # apply_async 可能已送达但客户端收到异常；worker 已领取时不能释放冻结。
        return
    error_message = "视频提词任务的 Redis 注册和 Celery 投递均失败，请重新执行步骤2"
    await mark_video_prompt_dispatch_failed_v2(
        db,
        config=SHOT_REPLICATE_V2,
        project_id=project_id,
        step_id=step_id,
        error_message=error_message,
    )
    log_celery_dispatch_compensated(billing_context, error=error_message)
    raise HTTPException(status_code=503, detail=error_message)


@router.post(
    "/segments/{segment_id}/replication-projects",
    response_model=ShotReplicateActionOut,
)
async def create_project_v2(
    segment_id: str = Path(...),
    req: ShotReplicateProjectCreateV2 = Body(...),
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    ensure_v2_celery_enabled()
    try:
        segment = await get_segment_for_user(
            db, segment_id=segment_id, user=current_user, for_update=True
        )
        result = await create_shot_replicate_project_v2(
            db, current_user=current_user, segment=segment, req=req
        )
        project_id = str(result.project.id)
        step_id = str(result.prompt_step.id)
        created_new = bool(result.created_new)
        billing_context = _dispatch_context(
            user_id=str(result.project.user_id),
            project_id=project_id,
            step_id=step_id,
            step_version=int(result.prompt_step.version or 1),
        )
        await db.commit()
    except IntegrityError as exc:
        await db.rollback()
        if not req.idempotency_key or not is_project_idempotency_conflict(exc):
            raise HTTPException(status_code=500, detail="项目创建失败") from exc
        segment = await get_segment_for_user(
            db, segment_id=segment_id, user=current_user, for_update=True
        )
        result = await create_shot_replicate_project_v2(
            db, current_user=current_user, segment=segment, req=req
        )
        project_id = str(result.project.id)
        step_id = str(result.prompt_step.id)
        created_new = bool(result.created_new)
        billing_context = _dispatch_context(
            user_id=str(result.project.user_id),
            project_id=project_id,
            step_id=step_id,
            step_version=int(result.prompt_step.version or 1),
        )
        await db.commit()
    except HTTPException:
        await db.rollback()
        raise
    except Exception as exc:
        await db.rollback()
        raise HTTPException(status_code=500, detail="创建拆镜复刻 V2 项目失败") from exc

    if created_new:
        await _dispatch_or_mark_failed(
            db, project_id=project_id, step_id=step_id, billing_context=billing_context
        )
    return ShotReplicateActionOut(
        message="V2 项目已创建，视频提词已自动提交" if created_new else "已返回现有幂等项目",
        project_id=project_id,
        step_id=step_id,
        detail=await _detail(db, current_user, project_id),
    )


@router.get("/projects/{project_id}", response_model=ShotReplicateTaskDetailOut)
async def get_project_v2(
    project_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    return await _detail(db, current_user, project_id)


@router.post(
    "/projects/{project_id}/steps/{step_id}/retry-video-prompt",
    response_model=ShotReplicateActionOut,
)
async def retry_video_prompt_v2(
    project_id: str,
    step_id: str,
    req: ModuleVideoPromptRetryV2,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    ensure_v2_celery_enabled()
    try:
        project, new_step = await rebuild_video_prompt_step_v2(
            db,
            config=SHOT_REPLICATE_V2,
            current_user=current_user,
            project_id=project_id,
            source_prompt_step_id=step_id,
            video_config=req.video_config,
        )
        project_id_value = str(project.id)
        step_id_value = str(new_step.id)
        billing_context = _dispatch_context(
            user_id=str(project.user_id),
            project_id=project_id_value,
            step_id=step_id_value,
            step_version=int(new_step.version or 1),
        )
        await db.commit()
    except HTTPException:
        await db.rollback()
        raise
    await _dispatch_or_mark_failed(
        db,
        project_id=project_id_value,
        step_id=step_id_value,
        billing_context=billing_context,
    )
    return ShotReplicateActionOut(
        message="视频提词已重新提交",
        project_id=project_id_value,
        step_id=step_id_value,
        detail=await _detail(db, current_user, project_id_value),
    )


@router.put(
    "/projects/{project_id}/steps/{step_id}/video-prompt-schema",
    response_model=ShotReplicateActionOut,
)
async def update_video_prompt_schema_route_v2(
    project_id: str,
    step_id: str,
    req: ModuleVideoPromptSchemaUpdateV2,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    try:
        project, step = await update_video_prompt_schema_v2(
            db,
            config=SHOT_REPLICATE_V2,
            current_user=current_user,
            project_id=project_id,
            step_id=step_id,
            req=req,
        )
        project_id_value = str(project.id)
        step_id_value = str(step.id)
        await db.commit()
    except HTTPException:
        await db.rollback()
        raise
    return ShotReplicateActionOut(
        message="视频提词已保存",
        project_id=project_id_value,
        step_id=step_id_value,
        detail=await _detail(db, current_user, project_id_value),
    )


@router.post(
    "/projects/{project_id}/steps/{step_id}/generate-video",
    response_model=ShotReplicateActionOut,
)
async def generate_video_v2(
    project_id: str,
    step_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    ensure_v2_celery_enabled()
    try:
        project, step, task = await generate_video_from_prompt_v2(
            db,
            config=SHOT_REPLICATE_V2,
            current_user=current_user,
            project_id=project_id,
            prompt_step_id=step_id,
        )
        project_id_value = str(project.id)
        step_id_value = str(step.id)
        task_id = str(task.id)
        await db.commit()
    except HTTPException:
        await db.rollback()
        raise

    queued_task = await db.get(ChatGenerationTask, task_id)
    if queued_task is None:
        raise HTTPException(status_code=500, detail="视频生成任务提交后无法重新读取")
    try:
        await enqueue_generation_create(queued_task, reason="shot_replicate_v2_generate_video")
    except Exception as exc:
        raise HTTPException(status_code=503, detail="视频生成任务暂未投递，将由恢复任务自动补投") from exc

    return ShotReplicateActionOut(
        message="视频生成任务已提交",
        project_id=project_id_value,
        step_id=step_id_value,
        detail=await _detail(db, current_user, project_id_value),
    )


@router.delete("/projects/{project_id}", response_model=ShotReplicateDeleteOut)
async def delete_project_route_v2(
    project_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    try:
        payload = await delete_project_v2(
            db,
            config=SHOT_REPLICATE_V2,
            current_user=current_user,
            project_id=project_id,
        )
        pending_ids = list(payload.get("pending_delete_resource_ids") or [])
        await db.commit()
    except HTTPException:
        await db.rollback()
        raise
    if pending_ids:
        try:
            await cleanup_upload_resource_files_after_commit(db, resource_ids=pending_ids)
            await db.commit()
        except Exception:
            await db.rollback()
    return ShotReplicateDeleteOut(**payload)