"""add llm_usage_event for billing

Revision ID: 20260803_10
Revises: 20260803_09
Create Date: 2026-08-03
"""
from __future__ import annotations

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

revision: str = "20260803_10"
down_revision: str | None = "20260803_09"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    op.create_table(
        "llm_usage_event",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column(
            "biz_dt",
            sa.String(length=8),
            nullable=False,
            comment="业务日 YYYYMMDD（Asia/Shanghai）",
        ),
        sa.Column(
            "provider",
            sa.String(length=32),
            nullable=False,
            server_default="openrouter",
            comment="提供商 openrouter / dashscope 等",
        ),
        sa.Column("model", sa.String(length=128), nullable=True, comment="模型名"),
        sa.Column("agent_name", sa.String(length=128), nullable=True, comment="Agent 名称"),
        sa.Column("run_id", sa.String(length=128), nullable=True, comment="Agent run id"),
        sa.Column("iteration", sa.Integer(), nullable=True, comment="ReAct 迭代序号"),
        sa.Column(
            "prompt_tokens",
            sa.Integer(),
            nullable=False,
            server_default="0",
            comment="输入 token 数",
        ),
        sa.Column(
            "completion_tokens",
            sa.Integer(),
            nullable=False,
            server_default="0",
            comment="输出 token 数",
        ),
        sa.Column(
            "total_tokens",
            sa.Integer(),
            nullable=False,
            server_default="0",
            comment="总 token 数",
        ),
        sa.Column(
            "cost_usd",
            sa.Numeric(precision=14, scale=8),
            nullable=True,
            comment="本次调用费用（USD，优先取 provider 返回的 usage.cost）",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
            comment="记录时间",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_llm_usage_biz_dt", "llm_usage_event", ["biz_dt"])
    op.create_index("idx_llm_usage_biz_agent", "llm_usage_event", ["biz_dt", "agent_name"])
    op.create_index("idx_llm_usage_biz_model", "llm_usage_event", ["biz_dt", "model"])
    op.create_index("idx_llm_usage_run_id", "llm_usage_event", ["run_id"])
    op.create_index("idx_llm_usage_created", "llm_usage_event", ["created_at"])


def downgrade() -> None:
    op.drop_index("idx_llm_usage_created", table_name="llm_usage_event")
    op.drop_index("idx_llm_usage_run_id", table_name="llm_usage_event")
    op.drop_index("idx_llm_usage_biz_model", table_name="llm_usage_event")
    op.drop_index("idx_llm_usage_biz_agent", table_name="llm_usage_event")
    op.drop_index("idx_llm_usage_biz_dt", table_name="llm_usage_event")
    op.drop_table("llm_usage_event")
