from __future__ import annotations

from datetime import datetime
from typing import Any

from sqlalchemy import Boolean, DateTime, Float, ForeignKey, Index, Integer, JSON, String, Text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column

from app.models.base import Base, SoftDeleteMixin, TimestampMixin

_STEP_JSON_TYPE = JSON().with_variant(JSONB, "postgresql")


class ModuleGenerationStep(Base, TimestampMixin, SoftDeleteMixin):
    """通用模块生成步骤表。

    V1 固定五步：material_input / image_prompt_optimize / image_generate /
    video_prompt_optimize / video_generate。

    V2 固定三步：material_input / video_prompt_optimize / video_generate。
    version 只表示同一步骤的重建版本，不表示项目流程版本；项目版本由
    ModuleGenerationProject.flow_version 保存。

    input_json / output_json 使用 JSON/JSONB 存储。
    建议结构：
    input_json = {
        "schema_version": "hot_opening_step_io_v1",
        "step_code": "...",
        "source": {...},
        "payload": {...},
        "context": {...}
    }
    output_json = {
        "schema_version": "hot_opening_step_io_v1",
        "step_code": "...",
        "status": "completed|failed|...",
        "payload": {...},
        "result": {...},
        "usage": {...},
        "error": {...}
    }
    """

    __tablename__ = "module_generation_steps"
    __table_args__ = (
        Index("idx_module_generation_steps_project_current", "project_id", "is_current", "deleted_at"),
        Index("idx_module_generation_steps_project_code", "project_id", "step_code", "is_current"),
        Index("idx_module_generation_steps_chat_task", "chat_task_id"),
    )

    id: Mapped[str] = mapped_column(String(32), primary_key=True)
    project_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("module_generation_projects.id", ondelete="CASCADE"), index=True, nullable=False
    )
    user_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="CASCADE"), index=True, nullable=False
    )
    module: Mapped[str] = mapped_column(String(64), index=True, nullable=False)
    step_index: Mapped[int] = mapped_column(Integer, index=True, nullable=False)
    step_code: Mapped[str] = mapped_column(String(64), index=True, nullable=False)
    status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
    version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)
    is_current: Mapped[bool] = mapped_column(Boolean, default=True, index=True)
    parent_step_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
    source_step_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
    chat_task_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("chat_generation_tasks.id", ondelete="SET NULL"), nullable=True, index=True
    )
    input_json: Mapped[dict[str, Any] | list[Any] | None] = mapped_column(_STEP_JSON_TYPE, nullable=True)
    output_json: Mapped[dict[str, Any] | list[Any] | None] = mapped_column(_STEP_JSON_TYPE, nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    # 提词优化成本快照，避免后台/详情页反复解析 output_json.usage。
    token_usage_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
    model_config_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
    input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
    output_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
    total_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
    text_credits_cost: Mapped[float | None] = mapped_column(Float, nullable=True)