"""
TokenUsage model for tracking OpenAI API token consumption
"""

from sqlalchemy import Column, String, Integer, Numeric, DateTime, JSON, ForeignKey
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
import uuid

from app.database import Base


class TokenUsage(Base):
    """
    TokenUsage model for tracking and analyzing token consumption and costs
    """

    __tablename__ = "token_usage"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
    user_id = Column(
        UUID(as_uuid=True),
        ForeignKey("users.id", ondelete="CASCADE"),
        nullable=False,
        index=True
    )
    conversation_id = Column(
        UUID(as_uuid=True),
        ForeignKey("conversations.id", ondelete="SET NULL"),
        index=True
    )
    message_id = Column(
        UUID(as_uuid=True),
        ForeignKey("messages.id", ondelete="SET NULL"),
        index=True
    )

    # Token counts
    prompt_tokens = Column(Integer, default=0, nullable=False)
    cached_tokens = Column(Integer, default=0, nullable=False)  # Cached input tokens (charged at lower rate)
    completion_tokens = Column(Integer, default=0, nullable=False)
    total_tokens = Column(Integer, default=0, nullable=False)

    # Model and cost information
    model = Column(String(100), nullable=False)  # e.g., gpt-5-nano-2025-08-07
    cost_usd = Column(Numeric(10, 6), default=0.0, nullable=False)  # Cost in USD

    # Operation type (chat, search, etc.)
    operation_type = Column(String(50))

    # Timestamp
    created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False, index=True)

    # JSON field for additional metadata
    meta_data = Column(JSON, default=dict, nullable=False)

    # Relationships
    user = relationship("User", back_populates="token_usage")
    conversation = relationship("Conversation", back_populates="token_usage")
    message = relationship("Message", back_populates="token_usage_records")

    def __repr__(self):
        return f"<TokenUsage {self.id} - {self.total_tokens} tokens>"