app/db/models.py

"""
Database models for billing and usage tracking.
"""
from sqlalchemy import Column, String, Integer, Float, DateTime, Boolean, ForeignKey, Text
from sqlalchemy.orm import relationship
from datetime import datetime
from typing import Optional

from app.db.database import Base


class Tenant(Base):
    """Tenant/organization model."""
    __tablename__ = "tenants"
    
    id = Column(String, primary_key=True, index=True)
    name = Column(String, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # Relationships
    plan = relationship("TenantPlan", back_populates="tenant", uselist=False)
    usage_events = relationship("UsageEvent", back_populates="tenant")
    daily_usage = relationship("UsageDaily", back_populates="tenant")
    monthly_usage = relationship("UsageMonthly", back_populates="tenant")


class TenantPlan(Base):
    """Tenant subscription plan."""
    __tablename__ = "tenant_plans"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    tenant_id = Column(String, ForeignKey("tenants.id"), unique=True, nullable=False, index=True)
    plan_name = Column(String, nullable=False, index=True)  # "starter", "growth", "pro"
    monthly_chat_limit = Column(Integer, nullable=False)  # -1 for unlimited
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # Relationships
    tenant = relationship("Tenant", back_populates="plan")


class UsageEvent(Base):
    """Individual usage event (each /chat request)."""
    __tablename__ = "usage_events"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    request_id = Column(String, unique=True, nullable=False, index=True)
    tenant_id = Column(String, ForeignKey("tenants.id"), nullable=False, index=True)
    user_id = Column(String, nullable=False, index=True)
    kb_id = Column(String, nullable=False)
    
    # LLM details
    provider = Column(String, nullable=False)  # "gemini" or "openai"
    model = Column(String, nullable=False)
    
    # Token usage
    prompt_tokens = Column(Integer, nullable=False, default=0)
    completion_tokens = Column(Integer, nullable=False, default=0)
    total_tokens = Column(Integer, nullable=False, default=0)
    
    # Cost tracking
    estimated_cost_usd = Column(Float, nullable=False, default=0.0)
    
    # Timestamp
    request_timestamp = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)
    
    # Relationships
    tenant = relationship("Tenant", back_populates="usage_events")


class UsageDaily(Base):
    """Daily aggregated usage per tenant."""
    __tablename__ = "usage_daily"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    tenant_id = Column(String, ForeignKey("tenants.id"), nullable=False, index=True)
    date = Column(DateTime, nullable=False, index=True)
    
    # Aggregated metrics
    total_requests = Column(Integer, nullable=False, default=0)
    total_tokens = Column(Integer, nullable=False, default=0)
    total_cost_usd = Column(Float, nullable=False, default=0.0)
    
    # Provider breakdown
    gemini_requests = Column(Integer, nullable=False, default=0)
    openai_requests = Column(Integer, nullable=False, default=0)
    
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # Unique constraint: one record per tenant per day
    __table_args__ = (
        {"sqlite_autoincrement": True},
    )
    
    # Relationships
    tenant = relationship("Tenant", back_populates="daily_usage")


class UsageMonthly(Base):
    """Monthly aggregated usage per tenant."""
    __tablename__ = "usage_monthly"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    tenant_id = Column(String, ForeignKey("tenants.id"), nullable=False, index=True)
    year = Column(Integer, nullable=False, index=True)
    month = Column(Integer, nullable=False, index=True)  # 1-12
    
    # Aggregated metrics
    total_requests = Column(Integer, nullable=False, default=0)
    total_tokens = Column(Integer, nullable=False, default=0)
    total_cost_usd = Column(Float, nullable=False, default=0.0)
    
    # Provider breakdown
    gemini_requests = Column(Integer, nullable=False, default=0)
    openai_requests = Column(Integer, nullable=False, default=0)
    
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # Unique constraint: one record per tenant per month
    __table_args__ = (
        {"sqlite_autoincrement": True},
    )
    
    # Relationships
    tenant = relationship("Tenant", back_populates="monthly_usage")