src/db/models.py

"""
Database Models and Session Management
"""

from datetime import datetime
from typing import Optional
from sqlalchemy import (
    Boolean, Column, DateTime, Float, Integer, String, Text, JSON, ForeignKey, Index
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, Session
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool

from src.core.config import settings

# Create declarative base
Base = declarative_base()


# Database Models
class User(Base):
    """User model for authentication"""
    __tablename__ = "users"
    
    id = Column(Integer, primary_key=True, index=True)
    email = Column(String(255), unique=True, index=True, nullable=False)
    hashed_password = Column(String(255), nullable=False)
    full_name = Column(String(255))
    is_active = Column(Boolean, default=True)
    is_superuser = Column(Boolean, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # Relationships
    api_keys = relationship("APIKey", back_populates="user", cascade="all, delete-orphan")
    requests = relationship("RequestLog", back_populates="user", cascade="all, delete-orphan")


class APIKey(Base):
    """API Key model for API authentication"""
    __tablename__ = "api_keys"
    
    id = Column(Integer, primary_key=True, index=True)
    key = Column(String(64), unique=True, index=True, nullable=False)
    name = Column(String(255))
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
    is_active = Column(Boolean, default=True)
    rate_limit_per_minute = Column(Integer, default=60)
    rate_limit_per_hour = Column(Integer, default=1000)
    created_at = Column(DateTime, default=datetime.utcnow)
    last_used_at = Column(DateTime)
    expires_at = Column(DateTime)
    
    # Relationships
    user = relationship("User", back_populates="api_keys")
    
    # Indexes
    __table_args__ = (
        Index('idx_apikey_user_active', 'user_id', 'is_active'),
    )


class RequestLog(Base):
    """Request logging for analytics and debugging"""
    __tablename__ = "request_logs"
    
    id = Column(Integer, primary_key=True, index=True)
    request_id = Column(String(64), unique=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=True)
    api_key_id = Column(Integer, ForeignKey("api_keys.id"), nullable=True)
    
    # Request details
    method = Column(String(10))
    path = Column(String(500))
    query_params = Column(JSON)
    status_code = Column(Integer)
    
    # Performance
    duration_ms = Column(Float)
    
    # Client info
    ip_address = Column(String(45))
    user_agent = Column(Text)
    
    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    
    # Relationships
    user = relationship("User", back_populates="requests")
    
    # Indexes
    __table_args__ = (
        Index('idx_request_user_created', 'user_id', 'created_at'),
        Index('idx_request_created', 'created_at'),
    )


class PredictionLog(Base):
    """ML prediction logging for analytics"""
    __tablename__ = "prediction_logs"
    
    id = Column(Integer, primary_key=True, index=True)
    request_id = Column(String(64), index=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=True)
    
    # Prediction details
    model_type = Column(String(50), index=True)  # deepfake, ai_text, anomaly
    input_type = Column(String(20))  # text, image, video, audio
    input_size = Column(Integer)  # bytes or character count
    
    # Results
    prediction = Column(String(50))
    confidence = Column(Float)
    details = Column(JSON)
    
    # Performance
    duration_ms = Column(Float)
    cached = Column(Boolean, default=False)
    
    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    
    # Indexes
    __table_args__ = (
        Index('idx_prediction_model_created', 'model_type', 'created_at'),
        Index('idx_prediction_user_created', 'user_id', 'created_at'),
    )


class SystemMetric(Base):
    """System performance metrics"""
    __tablename__ = "system_metrics"
    
    id = Column(Integer, primary_key=True, index=True)
    metric_name = Column(String(100), index=True)
    metric_value = Column(Float)
    labels = Column(JSON)
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    
    # Indexes
    __table_args__ = (
        Index('idx_metric_name_created', 'metric_name', 'created_at'),
    )


# Database Engine and Session
engine = create_engine(
    settings.DATABASE_URL,
    poolclass=QueuePool,
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,
    echo=settings.DEBUG
)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


# Dependency for FastAPI
def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# Database initialization
def create_tables():
    """Create all tables"""
    Base.metadata.create_all(bind=engine)


def drop_tables():
    """Drop all tables (use with caution!)"""
    Base.metadata.drop_all(bind=engine)


if __name__ == "__main__":
    print("Creating database tables...")
    create_tables()
    print("Tables created successfully!")