src/models/database.py

"""SQLAlchemy ORM models for AI Media OS."""

from sqlalchemy import Column, Integer, String, Text, Float, DateTime, Boolean, JSON, ForeignKey, Index, func, UniqueConstraint, Date
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship
from datetime import datetime
try:
    from pgvector.sqlalchemy import Vector
    _HAS_PGVECTOR = True
except ImportError:
    Vector = None
    _HAS_PGVECTOR = False

Base = declarative_base()


class Agent(Base):
    """AI agent persona configuration."""
    __tablename__ = "agents"

    id = Column(Integer, primary_key=True)
    name = Column(String, unique=True, nullable=False)
    description = Column(Text)
    current_version_id = Column(Integer, ForeignKey("agent_versions.id"))
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    versions = relationship("AgentVersion", back_populates="agent", cascade="all, delete-orphan", foreign_keys="[AgentVersion.agent_id]")
    current_version = relationship("AgentVersion", foreign_keys="[Agent.current_version_id]")


class AgentVersion(Base):
    """Versioned agent configuration."""
    __tablename__ = "agent_versions"
    __table_args__ = (
        UniqueConstraint("agent_id", "version", name="uq_agent_version"),
    )

    id = Column(Integer, primary_key=True)
    agent_id = Column(Integer, ForeignKey("agents.id", ondelete="CASCADE"), nullable=False)
    version = Column(Integer, nullable=False)
    config = Column(JSON, nullable=False)  # Model, temperature, etc.
    created_at = Column(DateTime, default=datetime.utcnow)
    created_by = Column(String)

    agent = relationship("Agent", back_populates="versions", foreign_keys="[AgentVersion.agent_id]")
    prompts = relationship("PromptTemplate", back_populates="agent_version", cascade="all, delete-orphan")
    posts = relationship("Post", back_populates="agent_version")


class PromptTemplate(Base):
    """Prompt templates for agent versions."""
    __tablename__ = "prompt_templates"

    id = Column(Integer, primary_key=True)
    agent_version_id = Column(Integer, ForeignKey("agent_versions.id", ondelete="CASCADE"), nullable=False)
    name = Column(String, nullable=False)
    prompt_text = Column(Text, nullable=False)
    description = Column(Text)
    created_at = Column(DateTime, default=datetime.utcnow)

    agent_version = relationship("AgentVersion", back_populates="prompts")


class Trend(Base):
    """Detected social trend."""
    __tablename__ = "trends"
    __table_args__ = (
        Index("ix_trends_status", "status"),
        Index("ix_trends_created_at", "created_at"),
        Index("ix_trends_source", "source"),
    )

    id = Column(Integer, primary_key=True)
    topic = Column(String, nullable=False)
    source = Column(String, nullable=False)
    source_url = Column(String)
    score = Column(Float)
    status = Column(String, default="pending_generation")
    used_for_post = Column(Boolean, default=False)
    raw_data = Column(JSON)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    topic_embedding = Column(Vector(1536) if _HAS_PGVECTOR else Text, nullable=True)

    posts = relationship("Post", back_populates="trend")


class MusicTrack(Base):
    """Royalty-free music tracks curated for social media posts."""
    __tablename__ = "music_tracks"

    id = Column(Integer, primary_key=True)
    title = Column(String, nullable=False)
    artist = Column(String)
    genre = Column(String)   # electronic|acoustic|pop|jazz|ambient|cinematic
    mood = Column(String)    # energetic|calm|upbeat|inspirational|modern|warm|romantic|sophisticated|playful
    bpm = Column(Integer)
    duration_secs = Column(Integer)
    preview_url = Column(String)
    cover_url = Column(String)
    tags = Column(JSON, default=list)
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)

    posts = relationship("Post", back_populates="music_track")


class Post(Base):
    """Generated social media post."""
    __tablename__ = "posts"
    __table_args__ = (
        Index("ix_posts_status", "status"),
        Index("ix_posts_approval_status", "approval_status"),
        Index("ix_posts_platform", "platform"),
        Index("ix_posts_published_at", "published_at"),
        Index("ix_posts_created_at", "created_at"),
    )

    id = Column(Integer, primary_key=True)
    trend_id = Column(Integer, ForeignKey("trends.id", ondelete="SET NULL"))
    agent_version_id = Column(Integer, ForeignKey("agent_versions.id"), nullable=True)
    content = Column(Text, nullable=False)
    post_type = Column(String, default="image")  # image | text_only | carousel | video | link | text_image
    image_url = Column(String)
    carousel_urls = Column(JSON)  # list of image URLs for carousel posts
    video_url = Column(String)
    link_url = Column(String)
    link_title = Column(String)
    link_description = Column(String)
    status = Column(String, default="draft")
    platform = Column(String)
    platform_post_id = Column(String)
    music_id = Column(Integer, ForeignKey("music_tracks.id", ondelete="SET NULL"), nullable=True)
    scheduled_at = Column(DateTime)
    published_at = Column(DateTime)
    
    approval_status = Column(String, default="pending")
    approval_notes = Column(Text)
    approved_by = Column(String)
    approved_at = Column(DateTime)
    
    token_count = Column(Integer)
    generation_cost = Column(Float)
    hallucination_score = Column(Float)
    quality_score = Column(Integer)
    quality_breakdown = Column(JSON)

    # Post-publish engagement metrics (fetched from Instagram ~24h after publish)
    ig_likes = Column(Integer)
    ig_comments = Column(Integer)
    ig_reach = Column(Integer)
    ig_saves = Column(Integer)
    ig_shares = Column(Integer)
    engagement_fetched_at = Column(DateTime)
    generation_error = Column(Text)
    content_embedding = Column(Vector(1536) if _HAS_PGVECTOR else Text, nullable=True)

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

    trend = relationship("Trend", back_populates="posts")
    agent_version = relationship("AgentVersion", back_populates="posts")
    media = relationship("Media", back_populates="post", cascade="all, delete-orphan")
    music_track = relationship("MusicTrack", back_populates="posts")


class Media(Base):
    """Media (image/video) associated with posts."""
    __tablename__ = "media"
    __table_args__ = (
        Index("ix_media_post_id", "post_id"),
        Index("ix_media_moderation_status", "moderation_status"),
    )

    id = Column(Integer, primary_key=True)
    post_id = Column(Integer, ForeignKey("posts.id", ondelete="CASCADE"))
    media_type = Column(String, nullable=False)
    file_name = Column(String, nullable=False)
    file_size = Column(Integer)
    mime_type = Column(String)
    r2_key = Column(String)
    r2_url = Column(String)
    
    moderation_status = Column(String, default="pending")
    nsfw_score = Column(Float)
    moderation_details = Column(JSON)
    thumbnail_url = Column(String)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    post = relationship("Post", back_populates="media")


class WorkflowExecution(Base):
    """Temporal workflow execution tracking."""
    __tablename__ = "workflow_executions"
    __table_args__ = (
        Index("ix_workflow_status", "status"),
        Index("ix_workflow_name", "workflow_name"),
        Index("ix_workflow_started_at", "started_at"),
        Index("ix_workflow_post_id", "post_id"),
        UniqueConstraint("workflow_id", name="uq_workflow_id"),
    )

    id = Column(Integer, primary_key=True)
    workflow_id = Column(String, nullable=False, unique=True)
    workflow_name = Column(String, nullable=False)
    trend_id = Column(Integer, ForeignKey("trends.id"))
    post_id = Column(Integer, ForeignKey("posts.id"))
    status = Column(String, default="running")
    started_at = Column(DateTime, default=datetime.utcnow)
    completed_at = Column(DateTime)
    error_message = Column(Text)
    execution_metadata = Column(JSON)


class ObservabilityTrace(Base):
    """Langfuse trace references and metadata."""
    __tablename__ = "observability_traces"
    __table_args__ = (
        Index("ix_trace_post_id", "post_id"),
        Index("ix_trace_workflow_id", "workflow_id"),
        Index("ix_trace_created_at", "created_at"),
    )

    id = Column(Integer, primary_key=True)
    trace_id = Column(String, unique=True)
    post_id = Column(Integer, ForeignKey("posts.id"))
    workflow_id = Column(String, ForeignKey("workflow_executions.workflow_id"))
    trace_name = Column(String)
    input = Column(JSON)
    output = Column(JSON)
    trace_metadata = Column(JSON)
    token_count = Column(Integer)
    cost_usd = Column(Float)
    duration_ms = Column(Integer)
    created_at = Column(DateTime, default=datetime.utcnow)


class ApiUsage(Base):
    """Track API usage and costs."""
    __tablename__ = "api_usage"
    __table_args__ = (
        UniqueConstraint("service", "period_start", "period_end", name="uq_service_period"),
    )

    id = Column(Integer, primary_key=True)
    service = Column(String, nullable=False)
    request_count = Column(Integer, default=0)
    token_count = Column(Integer, default=0)
    cost_usd = Column(Float, default=0.0)
    period_start = Column(DateTime)
    period_end = Column(DateTime)


class SystemAlert(Base):
    """System alerts and errors."""
    __tablename__ = "system_alerts"
    __table_args__ = (
        Index("ix_alert_severity", "severity"),
        Index("ix_alert_resolved", "resolved"),
        Index("ix_alert_created_at", "created_at"),
    )

    id = Column(Integer, primary_key=True)
    alert_type = Column(String, nullable=False)
    service = Column(String)
    message = Column(Text, nullable=False)
    severity = Column(String, default="warning")
    resolved = Column(Boolean, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    resolved_at = Column(DateTime)


class User(Base):
    """Web app user (authenticated via Firebase Google Auth)."""
    __tablename__ = "users"

    id = Column(Integer, primary_key=True)
    firebase_uid = Column(String, unique=True, nullable=False)
    email = Column(String, unique=True, nullable=False)
    display_name = Column(String)
    photo_url = Column(String)
    is_active = Column(Boolean, default=True)
    is_admin = Column(Boolean, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    last_login = Column(DateTime, default=datetime.utcnow)


class PlatformConfig(Base):
    """Per-platform configuration — tokens stored encrypted."""
    __tablename__ = "platform_configs"

    id = Column(Integer, primary_key=True)
    platform = Column(String, nullable=False, unique=True)   # instagram | twitter | linkedin | youtube
    is_enabled = Column(Boolean, default=False)
    access_token_enc = Column(Text)       # encrypted
    refresh_token_enc = Column(Text)      # encrypted
    account_id = Column(String)
    extra_config = Column(JSON, default={})
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class ModelConfig(Base):
    """LLM model configuration with fallback."""
    __tablename__ = "model_configs"

    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)           # e.g. "caption_generator"
    primary_model = Column(String, nullable=False)  # e.g. "google/gemma-3-27b-it:free"
    fallback_model = Column(String)                 # used when primary hits quota
    api_base = Column(String)
    api_key_enc = Column(Text)                      # encrypted
    is_active = Column(Boolean, default=True)
    last_error = Column(Text)
    last_error_at = Column(DateTime)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class BrandVoice(Base):
    """Active brand voice and prompt rules for content generation."""
    __tablename__ = "brand_voices"
    __table_args__ = (
        Index("ix_brand_voice_active", "is_active"),
    )

    id = Column(Integer, primary_key=True)
    brand_name = Column(String, default="Orbis")
    voice = Column(Text)
    audience = Column(Text)
    tone_rules = Column(Text)
    banned_words = Column(Text)
    cta_rules = Column(Text)
    hashtag_rules = Column(Text)
    auto_hashtags = Column(Text)   # space-separated tags always appended to every post
    caption_templates = Column(JSON, default=list)  # [{name, template}]
    example_posts = Column(Text)
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class BrandKnowledge(Base):
    """Brand knowledge base entries — FAQs, products, personas, context injected at generation time."""
    __tablename__ = "brand_knowledge"
    __table_args__ = (
        Index("ix_brand_knowledge_active", "is_active"),
        Index("ix_brand_knowledge_type", "content_type"),
    )

    id = Column(Integer, primary_key=True)
    title = Column(String, nullable=False)
    content = Column(Text, nullable=False)
    content_type = Column(String, default="general")  # general | faq | product | persona | competitor
    embedding = Column(Vector(1536) if _HAS_PGVECTOR else Text, nullable=True)
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class AgentFinding(Base):
    """AI agent team findings — daily analysis results."""
    __tablename__ = "agent_findings"
    __table_args__ = (
        Index("ix_agent_findings_run_date", "run_date"),
        Index("ix_agent_findings_agent_role", "agent_role"),
        Index("ix_agent_findings_severity", "severity"),
    )

    id = Column(Integer, primary_key=True)
    agent_role = Column(String, nullable=False)  # engineer|analyst|marketing|reviewer|architect|designer|manager
    category = Column(String)  # health|performance|content|quality|strategy|summary
    severity = Column(String, default="info")  # info|warning|critical
    title = Column(String, nullable=False)
    body = Column(Text)  # LLM-written natural language finding
    action_items = Column(JSON, default=list)  # list of strings
    data_snapshot = Column(JSON, default=dict)  # raw numbers used to produce this finding
    run_date = Column(Date, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)


class ProviderKey(Base):
    """Per-provider LLM API credentials — stored encrypted, manageable from Settings UI."""
    __tablename__ = "provider_keys"

    id = Column(Integer, primary_key=True)
    provider = Column(String, nullable=False, unique=True)  # openrouter|anthropic|groq|... or any custom slug
    display_name = Column(String)
    api_key_enc = Column(Text)      # Fernet-encrypted
    api_base = Column(String)       # required for custom providers, optional override for known ones
    is_enabled = Column(Boolean, default=True)
    models = Column(JSON, default=list)          # ordered model IDs to try
    disabled_models = Column(JSON, default=list) # model IDs toggled off by user
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class Notification(Base):
    """In-app notifications (token exhausted, publish failed, approval needed)."""
    __tablename__ = "notifications"
    __table_args__ = (
        Index("ix_notif_user_id", "user_id"),
        Index("ix_notif_read", "is_read"),
    )

    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey("users.id", ondelete="CASCADE"))
    type = Column(String, nullable=False)   # approval_needed | publish_failed | quota_exceeded
    title = Column(String, nullable=False)
    message = Column(Text)
    payload = Column(JSON, default={})      # e.g. {"workflow_id": "...", "post_id": 1}
    is_read = Column(Boolean, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)