models.py

"""
ORBIT – SQLAlchemy Models
Tables: User, UserSettings
"""
from datetime import datetime
from extensions import db


class User(db.Model):
    __tablename__ = "user"

    id         = db.Column(db.Integer, primary_key=True)
    google_id  = db.Column(db.String(128), unique=True, nullable=False, index=True)
    email      = db.Column(db.String(254), unique=True, nullable=False)
    name       = db.Column(db.String(200), nullable=False)
    picture    = db.Column(db.String(1000), default="")
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    last_login = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationship
    settings = db.relationship("UserSettings", back_populates="user", uselist=False, cascade="all, delete-orphan")
    sessions = db.relationship("ChatSession", back_populates="user", cascade="all, delete-orphan")

    def to_dict(self):
        return {
            "id":       self.id,
            "email":    self.email,
            "name":     self.name,
            "picture":  self.picture,
        }


class UserSettings(db.Model):
    __tablename__ = "user_settings"

    id               = db.Column(db.Integer, primary_key=True)
    user_id          = db.Column(db.Integer, db.ForeignKey("user.id"), unique=True, nullable=False)

    # Active provider & endpoint
    provider         = db.Column(db.String(64),  default="OpenRouter")
    base_url         = db.Column(db.String(512),  default="https://openrouter.ai/api/v1/chat/completions")
    api_key          = db.Column(db.String(512),  default="")
    current_model    = db.Column(db.String(256),  default="baidu/cobuddy:free")

    # Per-provider model libraries (JSON arrays)
    models_openrouter = db.Column(db.JSON, nullable=False, default=list)
    models_nvidia     = db.Column(db.JSON, nullable=False, default=list)

    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationship
    user = db.relationship("User", back_populates="settings")

    def active_models(self):
        """Return the model list for the currently selected provider."""
        if self.provider == "Nvidia NIM":
            return self.models_nvidia or []
        return self.models_openrouter or []

    def to_dict(self, mask_key: bool = False):
        key = self.api_key or ""
        if mask_key and key:
            key = key[:8] + "…" + key[-4:]
        return {
            "provider":          self.provider,
            "base_url":          self.base_url,
            "api_key":           key,
            "current_model":     self.current_model,
            "models_openrouter": self.models_openrouter or [],
            "models_nvidia":     self.models_nvidia or [],
            "active_models":     self.active_models(),
        }


class ChatSession(db.Model):
    __tablename__ = "chat_session"

    id         = db.Column(db.String(36), primary_key=True)
    user_id    = db.Column(db.Integer, db.ForeignKey("user.id"), nullable=False)
    title      = db.Column(db.String(255), default="New Chat")
    messages   = db.Column(db.JSON, nullable=False, default=list)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationship
    user = db.relationship("User", back_populates="sessions")

    def to_dict(self):
        return {
            "id":         self.id,
            "title":      self.title,
            "messages":   self.messages or [],
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }