from __future__ import annotations

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

from fastapi import HTTPException
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.models.chat_generation_task import ChatGenerationTask
from app.models.user import User
from app.schemas.generation_ai import GenerationAIReference, GenerationAITaskCreate
from app.services.generation.ai.engine_service import (
    IMAGE_DEFAULT_PROPORTION,
    IMAGE_DEFAULT_PX,
    IMAGE_DEFAULT_SIZE,
    VIDEO_DEFAULT_RATIO,
    VIDEO_DEFAULT_RESOLUTION,
    build_image_snapshot as _build_image_snapshot,
    build_video_snapshot as _build_video_snapshot,
    get_image_engine as _get_image_engine,
    get_video_engine as _get_video_engine,
    image_supported_sizes as _image_supported_sizes,
    normalize_px,
    parse_json_list as _parse_list,
)
from app.services.generation.billing_service import OWNER_CHAT_GENERATION_TASK, charge_generation_media_by_params
from app.services.generation.media_reference_service import (
    calculate_media_reference_usage,
    validate_media_reference_usage_for_engine,
)
from app.services.resource_capacity_service import assert_user_resource_capacity_available
from app.services.video_upscale.snapshot_service import build_video_upscale_snapshot
from app.services.private_portrait.reference_resolver import resolve_private_portrait_references
from app.utils.id_gen import generate_id


def _json(data: Any) -> str | None:
    if data is None:
        return None
    return json.dumps(data, ensure_ascii=False, default=str)


def _build_backend_idempotency_key(*, generation_mode: str, gen_type: str, task_id: str) -> str:
    """模块生成关联 ChatGenerationTask 的幂等键由后端生成。

    不再接收前端透传，避免 user_id + generation_mode + idempotency_key
    唯一索引被前端固定 key 或重复 key 拦截。
    """
    return f"{generation_mode}:{gen_type}:{task_id}"[:64]


async def create_chat_generation_task_for_module(
    db: AsyncSession,
    *,
    current_user: User,
    generation_mode: str,
    gen_type: str,
    original_prompt: str,
    optimized_prompt: str | None = None,
    engine_id: str | None = None,
    media_references: list[dict[str, Any]] | None = None,
    idempotency_key: str | None = None,
    image_size: str | None = None,
    image_proportion: str | None = None,
    image_px: str | None = None,
    duration: int | None = None,
    aspect_ratio: str | None = None,
    resolution: str | None = None,
    billing_project_name: str = "模块生成任务",
    billing_description_prefix: str = "模块生成-",
    billing_source_module: str | None = None,
    billing_source_project_id: str | None = None,
    billing_source_step_id: str | None = None,
    billing_source_step_code: str | None = None,
    billing_scene: str | None = None,
) -> ChatGenerationTask:
    """创建可复用的 ChatGenerationTask 子任务。

    和 /generation-ai 普通任务不同，generation_mode 由业务模块传入，
    但仍复用同一套引擎校验、扣费、Celery 创建/轮询/下载逻辑。
    """
    gen_type = gen_type.lower().strip()
    if gen_type not in ("image", "video"):
        raise HTTPException(status_code=400, detail="gen_type 仅支持 image 或 video")

    task_id = generate_id()
    now = datetime.now(timezone.utc)
    refs = media_references or []
    refs = await resolve_private_portrait_references(
        db,
        user_id=current_user.id,
        media_references=refs,
        gen_type=gen_type,
    )
    backend_idempotency_key = _build_backend_idempotency_key(
        generation_mode=generation_mode,
        gen_type=gen_type,
        task_id=task_id,
    )

    await assert_user_resource_capacity_available(db, current_user.id)

    if gen_type == "image":
        engine = await _get_image_engine(db, engine_id)
        sizes = _image_supported_sizes(engine)
        size = image_size or engine.default_size or IMAGE_DEFAULT_SIZE
        proportion = image_proportion or IMAGE_DEFAULT_PROPORTION
        px = normalize_px(image_px)
        if sizes:
            if size not in sizes:
                raise HTTPException(status_code=400, detail=f"图片分辨率档位不支持: {size}")
            if proportion not in sizes.get(size, {}):
                raise HTTPException(status_code=400, detail=f"图片比例不支持: {proportion}")
            px = px or normalize_px((sizes.get(size) or {}).get(proportion))
        px = px or IMAGE_DEFAULT_PX
        reference_usage = calculate_media_reference_usage(refs, include=True)
        validate_media_reference_usage_for_engine(reference_usage, gen_type="image", engine=engine)
        media_billing = await charge_generation_media_by_params(
            db,
            user_id=current_user.id,
            record_id=task_id,
            gen_type="image",
            image_size=size,
            engine_id=engine.id,
            input_image_count=reference_usage.image_count or None,
            project_name=billing_project_name,
            description_prefix=billing_description_prefix,
            owner_type=OWNER_CHAT_GENERATION_TASK,
            attempt_no=1,
            source_module=billing_source_module,
            source_project_id=billing_source_project_id,
            source_step_id=billing_source_step_id,
            source_step_code=billing_source_step_code,
            billing_scene=billing_scene,
        )
        snapshot = _build_image_snapshot(engine, size, proportion, px)
        snapshot["generation_count"] = 1
        task = ChatGenerationTask(
            id=task_id,
            created_at=now,
            resource_generation_started_at=now,
            generation_attempt_no=1,
            user_id=current_user.id,
            original_prompt=original_prompt,
            optimized_prompt=optimized_prompt,
            gen_type="image",
            image_size=size,
            image_proportion=proportion,
            image_px=px,
            status="generating",
            generation_mode=generation_mode,
            pipeline_stage="queued",
            engine_id=engine.id,
            engine_snapshot_json=_json(snapshot),
            media_references=_json(refs) if refs else None,
            credits_cost=round(media_billing.total_charged, 2),
            idempotency_key=backend_idempotency_key,
            deadline_at=now + timedelta(minutes=settings.CHATAPI_ASYNC_IMAGE_DEADLINE_MINUTES),
        )
    else:
        engine = await _get_video_engine(db, engine_id)
        ratio = aspect_ratio or VIDEO_DEFAULT_RATIO
        selected_resolution = resolution or VIDEO_DEFAULT_RESOLUTION
        selected_duration = duration or 4
        ratios = _parse_list(engine.supported_ratios, [])
        resolutions = _parse_list(engine.supported_resolutions, [])
        durations = _parse_list(engine.supported_durations, [])
        if ratios and ratio not in ratios:
            raise HTTPException(status_code=400, detail=f"视频比例不支持: {ratio}")
        if resolutions and selected_resolution not in resolutions:
            raise HTTPException(status_code=400, detail=f"视频分辨率不支持: {selected_resolution}")
        if durations and selected_duration not in durations:
            raise HTTPException(status_code=400, detail=f"视频时长不支持: {selected_duration}")
        if engine.max_duration and selected_duration > engine.max_duration:
            raise HTTPException(status_code=400, detail=f"视频时长不能超过 {engine.max_duration} 秒")
        provider_generation_resolution, upscale_enabled_snapshot, upscale_snapshot_json = await build_video_upscale_snapshot(
            db,
            target_resolution=selected_resolution,
            aspect_ratio=ratio,
            supported_provider_resolutions=resolutions,
        )
        reference_usage = calculate_media_reference_usage(refs, include=True)
        validate_media_reference_usage_for_engine(reference_usage, gen_type="video", engine=engine)
        media_billing = await charge_generation_media_by_params(
            db,
            user_id=current_user.id,
            record_id=task_id,
            gen_type="video",
            duration=selected_duration,
            resolution=selected_resolution,
            engine_id=engine.id,
            input_video_duration=reference_usage.input_video_duration or None,
            input_image_count=reference_usage.image_count or None,
            project_name=billing_project_name,
            description_prefix=billing_description_prefix,
            owner_type=OWNER_CHAT_GENERATION_TASK,
            attempt_no=1,
            source_module=billing_source_module,
            source_project_id=billing_source_project_id,
            source_step_id=billing_source_step_id,
            source_step_code=billing_source_step_code,
            billing_scene=billing_scene,
        )
        snapshot = _build_video_snapshot(engine, ratio, selected_resolution, selected_duration)
        snapshot["generation_count"] = 1
        snapshot["provider_generation_resolution"] = provider_generation_resolution
        snapshot["video_upscale_enabled_snapshot"] = upscale_enabled_snapshot
        task = ChatGenerationTask(
            id=task_id,
            created_at=now,
            resource_generation_started_at=now,
            generation_attempt_no=1,
            user_id=current_user.id,
            original_prompt=original_prompt,
            optimized_prompt=optimized_prompt,
            gen_type="video",
            duration=selected_duration,
            aspect_ratio=ratio,
            resolution=selected_resolution,
            provider_generation_resolution=provider_generation_resolution,
            video_upscale_enabled_snapshot=upscale_enabled_snapshot,
            video_upscale_snapshot_json=upscale_snapshot_json,
            image_size=image_size or IMAGE_DEFAULT_SIZE,
            image_proportion=image_proportion or IMAGE_DEFAULT_PROPORTION,
            image_px=normalize_px(image_px) or IMAGE_DEFAULT_PX,
            status="generating",
            generation_mode=generation_mode,
            pipeline_stage="queued",
            engine_id=engine.id,
            engine_snapshot_json=_json(snapshot),
            media_references=_json(refs) if refs else None,
            credits_cost=round(media_billing.total_charged, 2),
            idempotency_key=backend_idempotency_key,
            deadline_at=now + timedelta(hours=settings.CHATAPI_ASYNC_VIDEO_FINAL_DEADLINE_HOURS),
        )

    db.add(task)
    await db.flush()
    return task