from __future__ import annotations

from dataclasses import dataclass

from app.enums.hot_opening_replicate import (
    HotOpeningGenerationModeEnum,
    HotOpeningStepIOSchemaVersionEnum,
    ModuleCodeEnum as HotModuleCodeEnum,
)
from app.enums.shot_replicate import (
    ModuleCodeEnum as ShotModuleCodeEnum,
    ShotReplicateGenerationModeEnum,
    ShotReplicateStepIOSchemaVersionEnum,
)
from app.enums.module_generation_flow import ModuleGenerationFlowConfig

MATERIAL_INPUT = "material_input"
VIDEO_PROMPT_OPTIMIZE = "video_prompt_optimize"
VIDEO_GENERATE = "video_generate"


@dataclass(frozen=True, slots=True)
class ModuleGenerationV2Config:
    module: str
    generation_mode: str
    io_schema_version: str
    display_name: str
    project_not_found_message: str
    material_video_locked: bool

    @property
    def flow_config(self) -> ModuleGenerationFlowConfig:
        return ModuleGenerationFlowConfig(
            module=self.module,
            step_index_map={
                MATERIAL_INPUT: 1,
                VIDEO_PROMPT_OPTIMIZE: 2,
                VIDEO_GENERATE: 3,
            },
            material_step_code=MATERIAL_INPUT,
            image_prompt_step_code="__v2_no_image_prompt__",
            image_generate_step_code="__v2_no_image_generate__",
            video_prompt_step_code=VIDEO_PROMPT_OPTIMIZE,
            video_generate_step_code=VIDEO_GENERATE,
            project_not_found_message=self.project_not_found_message,
            step_not_found_message="V2 子任务不存在",
            cancel_chat_task_error_message=f"{self.display_name}步骤被重建或删除，旧生成任务已取消",
            material_video_url_editable=not self.material_video_locked,
            step_io_schema_version=self.io_schema_version,
            expected_flow_version="v2",
        )


HOT_OPENING_V2 = ModuleGenerationV2Config(
    module=HotModuleCodeEnum.HOT_OPENING_REPLICATE.value,
    generation_mode=HotOpeningGenerationModeEnum.HOT_OPENING_REPLICATE.value,
    io_schema_version=HotOpeningStepIOSchemaVersionEnum.V2.value,
    display_name="爆款开头复刻",
    project_not_found_message="爆款开头复刻 V2 项目不存在",
    material_video_locked=True,
)

SHOT_REPLICATE_V2 = ModuleGenerationV2Config(
    module=ShotModuleCodeEnum.SHOT_REPLICATE.value,
    generation_mode=ShotReplicateGenerationModeEnum.SHOT_REPLICATE.value,
    io_schema_version=ShotReplicateStepIOSchemaVersionEnum.V2.value,
    display_name="拆镜复刻",
    project_not_found_message="拆镜复刻 V2 项目不存在",
    material_video_locked=True,
)

CONFIG_BY_MODULE = {
    HOT_OPENING_V2.module: HOT_OPENING_V2,
    SHOT_REPLICATE_V2.module: SHOT_REPLICATE_V2,
}