"""003_business_tables

Revision ID: 003
Revises: 002
Create Date: 2026-06-09
Alembic 003: 业务主体7张表 + ai_call_logs
products / benchmark_notes / banned_words /
generation_tasks / text_candidates / image_candidates /
delivery_packages / ai_call_logs
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

revision: str = "003"
down_revision: Union[str, None] = "002"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

_TASK_STATUS = sa.Enum(
    "pending", "generating", "pending_selection",
    "pending_review", "approved", "rejected", "archived",
    name="taskstatus",
)
_REVIEW_STATUS = sa.Enum("pending", "approved", "rejected", name="reviewstatus")
_CANDIDATE_SOURCE = sa.Enum("ai", "import", name="candidatesource")
_BANNED_LEVEL = sa.Enum("auto_fix", "soft_warn", "hard_block", name="bannedwordlevel")
_BANNED_STATUS = sa.Enum("pass", "auto_fixed", "soft_warn", "hard_block", name="bannedwordstatus")
_IMAGE_ROLE = sa.Enum("hook", "pain", "proof", "quality", "credit", "convert", "main", name="imagerole")
_PKG_STATUS = sa.Enum("pending", "ready", "downloaded", name="packagestatus")
_PRODUCT_SOURCE = sa.Enum("preset", "custom", name="productsource")


def upgrade() -> None:
    op.create_table(
        "products",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("name", sa.String(128), nullable=False),
        sa.Column("category", sa.String(64), nullable=True),
        sa.Column("source", _PRODUCT_SOURCE, nullable=False, server_default="custom"),
        sa.Column("selling_points", sa.Text(), nullable=True),
        sa.Column("style_tone", sa.String(128), nullable=True),
        sa.Column("text_angles", sa.Text(), nullable=True),
        sa.Column("custom_prompt", sa.Text(), nullable=True),
        sa.Column("is_active", sa.Boolean(), nullable=False, server_default="1"),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["workspace_id"], ["workspaces.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_products_workspace_id", "products", ["workspace_id"])

    op.create_table(
        "benchmark_notes",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("product_id", sa.BigInteger(), nullable=False),
        sa.Column("screenshot_url", sa.String(512), nullable=True),
        sa.Column("highlights", sa.Text(), nullable=True),
        sa.Column("link_url", sa.String(512), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["workspace_id"], ["workspaces.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["product_id"], ["products.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_benchmark_notes_product_id", "benchmark_notes", ["product_id"])

    op.create_table(
        "banned_words",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("word", sa.String(64), nullable=False),
        sa.Column("level", _BANNED_LEVEL, nullable=False),
        sa.Column("replacement", sa.String(128), nullable=True),
        sa.Column("updatable", sa.Boolean(), nullable=False, server_default="1"),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["workspace_id"], ["workspaces.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_banned_words_workspace_id", "banned_words", ["workspace_id"])

    op.create_table(
        "generation_tasks",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("product_id", sa.BigInteger(), nullable=False),
        sa.Column("operator_id", sa.BigInteger(), nullable=False),
        sa.Column("theme", sa.String(256), nullable=True),
        sa.Column("text_count", sa.Integer(), nullable=False, server_default="5"),
        sa.Column("image_count", sa.Integer(), nullable=False, server_default="3"),
        sa.Column("track", sa.String(16), nullable=False, server_default="ai"),
        sa.Column("status", _TASK_STATUS, nullable=False, server_default="pending"),
        sa.Column("mongo_trace_id", sa.String(24), nullable=True),
        sa.Column("review_status", _REVIEW_STATUS, nullable=True),
        sa.Column("reviewer_id", sa.BigInteger(), nullable=True),
        sa.Column("reviewed_at", sa.DateTime(), nullable=True),
        sa.Column("reject_reason", sa.Text(), nullable=True),
        sa.Column("approved_at", sa.DateTime(), nullable=True),
        sa.Column("archived_at", sa.DateTime(), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["workspace_id"], ["workspaces.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["product_id"], ["products.id"]),
        sa.ForeignKeyConstraint(["operator_id"], ["users.id"]),
        sa.ForeignKeyConstraint(["reviewer_id"], ["users.id"]),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_generation_tasks_workspace_status", "generation_tasks", ["workspace_id", "status"])

    op.create_table(
        "text_candidates",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("task_id", sa.BigInteger(), nullable=False),
        sa.Column("source", _CANDIDATE_SOURCE, nullable=False, server_default="ai"),
        sa.Column("angle_label", sa.String(64), nullable=True),
        sa.Column("content", sa.Text(), nullable=True),
        sa.Column("score_json", sa.Text(), nullable=True),
        sa.Column("banned_word_status", _BANNED_STATUS, nullable=False, server_default="pass"),
        sa.Column("eval_score", sa.Float(), nullable=True),
        sa.Column("is_selected", sa.Boolean(), nullable=False, server_default="0"),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["task_id"], ["generation_tasks.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_text_candidates_task_id", "text_candidates", ["task_id"])

    op.create_table(
        "image_candidates",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("task_id", sa.BigInteger(), nullable=False),
        sa.Column("role", _IMAGE_ROLE, nullable=False, server_default="main"),
        sa.Column("url", sa.String(512), nullable=True),
        sa.Column("strategy", sa.String(4), nullable=True),
        sa.Column("seq", sa.Integer(), nullable=False, server_default="1"),
        sa.Column("is_selected", sa.Boolean(), nullable=False, server_default="0"),
        sa.Column("eval_score", sa.Float(), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["task_id"], ["generation_tasks.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_image_candidates_task_id", "image_candidates", ["task_id"])

    op.create_table(
        "delivery_packages",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("task_id", sa.BigInteger(), nullable=False),
        sa.Column("status", _PKG_STATUS, nullable=False, server_default="pending"),
        sa.Column("package_path", sa.String(512), nullable=True),
        sa.Column("download_url", sa.String(512), nullable=True),
        sa.Column("expires_at", sa.DateTime(), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["workspace_id"], ["workspaces.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["task_id"], ["generation_tasks.id"]),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_table(
        "ai_call_logs",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column("workspace_id", sa.BigInteger(), nullable=False),
        sa.Column("user_id", sa.BigInteger(), nullable=False),
        sa.Column("key_id", sa.BigInteger(), nullable=True),
        sa.Column("task_id", sa.BigInteger(), nullable=True),
        sa.Column("provider", sa.String(32), nullable=True),
        sa.Column("model", sa.String(64), nullable=True),
        sa.Column("call_type", sa.String(32), nullable=True),
        sa.Column("prompt_tokens", sa.Integer(), nullable=True),
        sa.Column("completion_tokens", sa.Integer(), nullable=True),
        sa.Column("success", sa.Boolean(), nullable=False, server_default="1"),
        sa.Column("error_code", sa.String(32), nullable=True),
        sa.Column("latency_ms", sa.Integer(), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["user_id"], ["users.id"]),
        sa.ForeignKeyConstraint(["key_id"], ["user_api_keys.id"]),
        sa.ForeignKeyConstraint(["task_id"], ["generation_tasks.id"]),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_ai_call_logs_workspace_user", "ai_call_logs", ["workspace_id", "user_id"])
    op.create_index("idx_ai_call_logs_task_id", "ai_call_logs", ["task_id"])


def downgrade() -> None:
    op.drop_index("idx_ai_call_logs_task_id", "ai_call_logs")
    op.drop_index("idx_ai_call_logs_workspace_user", "ai_call_logs")
    op.drop_table("ai_call_logs")
    op.drop_table("delivery_packages")
    op.drop_index("idx_image_candidates_task_id", "image_candidates")
    op.drop_table("image_candidates")
    op.drop_index("idx_text_candidates_task_id", "text_candidates")
    op.drop_table("text_candidates")
    op.drop_index("idx_generation_tasks_workspace_status", "generation_tasks")
    op.drop_table("generation_tasks")
    op.drop_index("idx_banned_words_workspace_id", "banned_words")
    op.drop_table("banned_words")
    op.drop_index("idx_benchmark_notes_product_id", "benchmark_notes")
    op.drop_table("benchmark_notes")
    op.drop_index("idx_products_workspace_id", "products")
    op.drop_table("products")
    # MySQL 不支持 DROP TYPE，枚举存 VARCHAR 无需清理