"""usage_events 表(v2 形态)+ messages.model_profile 列。

Revision ID: 0006
Revises: 0005
Create Date: 2026-05-19

模型切换 + 用量统计(grill 共识):
- task 默认模型走 tasks.model_profile(0001 起就存,本次开始真用)
- 每条 assistant message 标注实际用的模型 → messages.model_profile(本次加列)
- 用量按 message 记成一行 usage_event(JSONB polymorphic units),前期只 chat kind;
  未来扩 image/video/audio 不动 schema。0004 删掉的 v1 形态(id/user_id/task_id/run_id/
  kind/value/ts)字段不够,本次按 grill Q6 设计重建一张 v2 形态。
- task_id ondelete=CASCADE(简单,跟 messages 一致);message_id ondelete=SET NULL
  (单条 message 不会主动删,留作未来防御);user_id 仍是 NOT NULL —— 计费 / 限额查询
  按 user 维度直出,不靠 JOIN。
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision: str = "0006"
down_revision: Union[str, None] = "0005"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.add_column(
        "messages",
        sa.Column("model_profile", sa.Text(), nullable=True),
    )

    op.create_table(
        "usage_events",
        sa.Column("event_id", postgresql.UUID(as_uuid=True), primary_key=True,
                  server_default=sa.text("gen_random_uuid()")),
        sa.Column("user_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("users.user_id"), nullable=False),
        sa.Column("task_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("tasks.task_id", ondelete="CASCADE"), nullable=False),
        sa.Column("message_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("messages.message_id", ondelete="SET NULL"), nullable=True),
        sa.Column("kind", sa.Text(), nullable=False),
        sa.Column("model_profile", sa.Text(), nullable=False),
        sa.Column("units", postgresql.JSONB(), nullable=False),
        sa.Column("cost_usd", sa.Numeric(12, 6), nullable=False, server_default="0"),
        sa.Column("created_at", sa.DateTime(timezone=True),
                  server_default=sa.text("now()"), nullable=False),
    )
    op.create_index("ix_usage_user_created", "usage_events", ["user_id", "created_at"])
    op.create_index("ix_usage_task", "usage_events", ["task_id"])
    op.create_index("ix_usage_model_created", "usage_events", ["model_profile", "created_at"])


def downgrade() -> None:
    op.drop_index("ix_usage_model_created", table_name="usage_events")
    op.drop_index("ix_usage_task", table_name="usage_events")
    op.drop_index("ix_usage_user_created", table_name="usage_events")
    op.drop_table("usage_events")
    op.drop_column("messages", "model_profile")
