"""add ai cost tracking fields

Revision ID: 0018_ai_cost_tracking
Revises: 0017_share_drafts
"""

from alembic import context, op
import sqlalchemy as sa


revision = "0018_ai_cost_tracking"
down_revision = "0017_share_drafts"
branch_labels = None
depends_on = None


def upgrade() -> None:
    inspector = None if context.is_offline_mode() else sa.inspect(op.get_bind())
    model_columns = set() if inspector is None else {column["name"] for column in inspector.get_columns("sys_model")}
    _add_column_if_missing(model_columns, "sys_model", sa.Column("input_price_per_1k", sa.Numeric(12, 6), nullable=True))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("output_price_per_1k", sa.Numeric(12, 6), nullable=True))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("currency", sa.String(10), nullable=False, server_default="CNY"))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("usage_scenarios", sa.String(255), nullable=True))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_summary", sa.Integer(), nullable=False, server_default="1"))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_report", sa.Integer(), nullable=False, server_default="1"))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_fixed_info", sa.Integer(), nullable=False, server_default="1"))
    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_deep_chat", sa.Integer(), nullable=False, server_default="1"))

    log_columns = set() if inspector is None else {column["name"] for column in inspector.get_columns("sys_ai_request_log")}
    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("model_id", sa.BigInteger(), nullable=True))
    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("estimated_cost", sa.Numeric(18, 6), nullable=True))
    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("currency", sa.String(10), nullable=True))
    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("route_reason", sa.String(255), nullable=True))
    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("question_type", sa.String(50), nullable=True))
    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("knowledge_hit", sa.Integer(), nullable=False, server_default="0"))
    if "model_id" not in log_columns:
        op.create_index("ix_sys_ai_request_log_model_id", "sys_ai_request_log", ["model_id"])


def downgrade() -> None:
    inspector = None if context.is_offline_mode() else sa.inspect(op.get_bind())
    log_columns = {
        "model_id", "estimated_cost", "currency", "route_reason", "question_type", "knowledge_hit"
    } if inspector is None else {column["name"] for column in inspector.get_columns("sys_ai_request_log")}
    if "model_id" in log_columns:
        op.drop_index("ix_sys_ai_request_log_model_id", table_name="sys_ai_request_log")
    for column in ["knowledge_hit", "question_type", "route_reason", "currency", "estimated_cost", "model_id"]:
        if column in log_columns:
            op.drop_column("sys_ai_request_log", column)

    model_columns = {
        "input_price_per_1k", "output_price_per_1k", "currency", "usage_scenarios",
        "allow_summary", "allow_report", "allow_fixed_info", "allow_deep_chat",
    } if inspector is None else {column["name"] for column in inspector.get_columns("sys_model")}
    for column in [
        "allow_deep_chat",
        "allow_fixed_info",
        "allow_report",
        "allow_summary",
        "usage_scenarios",
        "currency",
        "output_price_per_1k",
        "input_price_per_1k",
    ]:
        if column in model_columns:
            op.drop_column("sys_model", column)


def _add_column_if_missing(columns: set[str], table: str, column: sa.Column) -> None:
    if column.name not in columns:
        op.add_column(table, column)
