from sqlalchemy import ForeignKey, Index, Integer, String
from sqlalchemy.orm import Mapped, mapped_column

from app.models.base import Base, TimestampMixin


class TokenUsage(Base, TimestampMixin):
    __tablename__ = "token_usage"
    __table_args__ = (
        Index("ix_token_usage_owner", "owner_type", "owner_id"),
        Index("ix_token_usage_biz_key", "biz_key"),
    )

    id: Mapped[str] = mapped_column(String(32), primary_key=True)
    model_config_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("model_configs.id", ondelete="SET NULL"), nullable=True
    )
    user_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="SET NULL"), nullable=True, index=True
    )
    input_tokens: Mapped[int] = mapped_column(Integer, default=0)
    output_tokens: Mapped[int] = mapped_column(Integer, default=0)
    total_tokens: Mapped[int] = mapped_column(Integer, default=0)

    owner_type: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
    owner_id: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
    biz_key: Mapped[str | None] = mapped_column(String(160), nullable=True, index=True)
    source_module: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
    source_step_code: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)