"""Nutzung je Modellaufruf und Modellpreise

Revision ID: 0008
Revises: 0007
Erstellt: 2026-10-22 09:41:27.635420
"""

import sqlalchemy as sa
from alembic import op


revision = '0008'
down_revision = '0007'
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.create_table('usage_records',
    sa.Column('id', sa.BigInteger(), sa.Identity(always=False), nullable=False),
    sa.Column('time', sa.DateTime(timezone=True), nullable=False),
    sa.Column('run_id', sa.Uuid(), nullable=True),
    sa.Column('team_id', sa.Integer(), nullable=True),
    sa.Column('workspace_id', sa.Integer(), nullable=True),
    sa.Column('user_id', sa.Integer(), nullable=True),
    sa.Column('model', sa.String(length=200), nullable=False),
    sa.Column('input_tokens', sa.Integer(), nullable=False),
    sa.Column('cached_tokens', sa.Integer(), nullable=False),
    sa.Column('output_tokens', sa.Integer(), nullable=False),
    sa.Column('reasoning_tokens', sa.Integer(), nullable=False),
    sa.Column('cost', sa.Numeric(precision=14, scale=6), nullable=True),
    sa.ForeignKeyConstraint(['team_id'], ['teams.id'], name=op.f('fk_usage_records_team_id_teams'), ondelete='SET NULL'),
    sa.ForeignKeyConstraint(['user_id'], ['users.id'], name=op.f('fk_usage_records_user_id_users'), ondelete='SET NULL'),
    sa.ForeignKeyConstraint(['workspace_id'], ['workspaces.id'], name=op.f('fk_usage_records_workspace_id_workspaces'), ondelete='SET NULL'),
    sa.PrimaryKeyConstraint('id', name=op.f('pk_usage_records'))
    )
    op.create_index(op.f('ix_usage_records_run_id'), 'usage_records', ['run_id'], unique=False)
    op.create_index(op.f('ix_usage_records_team_id'), 'usage_records', ['team_id'], unique=False)
    op.create_index(op.f('ix_usage_records_time'), 'usage_records', ['time'], unique=False)
    op.create_index(op.f('ix_usage_records_workspace_id'), 'usage_records', ['workspace_id'], unique=False)
    op.add_column('models', sa.Column('input_price', sa.Numeric(precision=12, scale=4), nullable=True))
    op.add_column('models', sa.Column('cached_price', sa.Numeric(precision=12, scale=4), nullable=True))
    op.add_column('models', sa.Column('output_price', sa.Numeric(precision=12, scale=4), nullable=True))
    # Bisherigen Verbrauch aus den Antworten übernehmen (je Antwort summiert, ohne Preis)
    op.execute("""
        INSERT INTO usage_records (time, run_id, team_id, workspace_id, user_id, model, input_tokens, cached_tokens,
                                   output_tokens, reasoning_tokens)
        SELECT m.created, m.run_id, w.team_id, c.workspace_id, r.created_by, m.model,
               coalesce((m.usage->>'input_tokens')::int, 0), 0, coalesce((m.usage->>'output_tokens')::int, 0), 0
        FROM messages m
        JOIN chat_sessions c ON c.id = m.chat_session_id
        JOIN workspaces w ON w.id = c.workspace_id
        LEFT JOIN task_runs r ON r.id = m.run_id
        WHERE m.role = 'assistant' AND m.usage IS NOT NULL AND (m.usage->>'total_tokens') IS NOT NULL
    """)


def downgrade() -> None:
    op.drop_column('models', 'output_price')
    op.drop_column('models', 'cached_price')
    op.drop_column('models', 'input_price')
    op.drop_index(op.f('ix_usage_records_workspace_id'), table_name='usage_records')
    op.drop_index(op.f('ix_usage_records_time'), table_name='usage_records')
    op.drop_index(op.f('ix_usage_records_team_id'), table_name='usage_records')
    op.drop_index(op.f('ix_usage_records_run_id'), table_name='usage_records')
    op.drop_table('usage_records')