src/models/database.py

"""
Database Models -- AI Reel Creator Platform
===========================================

SQLAlchemy async ORM models mapping exactly to configs/database_schema.sql.
"""

import uuid
from datetime import datetime
from typing import List, Optional

from sqlalchemy import (
    create_engine, Column, String, Text, Integer, BigInteger, Float,
    Boolean, DateTime, ForeignKey, ARRAY, JSON, func, event, select,
)
from sqlalchemy.orm import declarative_base, relationship, Session
from sqlalchemy.dialects.postgresql import UUID, JSONB, ENUM

_TRY_PGVECTOR = True
Vector = None

try:
    from pgvector.sqlalchemy import Vector as _Vector
    Vector = _Vector
except ImportError:
    _TRY_PGVECTOR = False
    class _VectorStub:
        def __init__(self, dim: int):
            self.dim = dim
    Vector = _VectorStub

Base = declarative_base()


class TimestampMixin:
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())


class Asset(Base, TimestampMixin):
    __tablename__ = "assets"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    file_path = Column(Text, nullable=False)
    file_name = Column(Text, nullable=False)
    asset_type = Column(ENUM("video", "image", "audio", name="asset_type_enum"), nullable=False)
    source = Column(Text)
    resolution = Column(Text)
    duration_ms = Column(Integer)
    frame_rate = Column(Float)
    file_size_bytes = Column(BigInteger)
    metadata = relationship("AssetMetadata", uselist=False, back_populates="asset")
    video_events = relationship("VideoEvent", back_populates="asset")


class AssetMetadata(Base, TimestampMixin):
    __tablename__ = "asset_metadata"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    asset_id = Column(UUID(as_uuid=True), ForeignKey("assets.id", ondelete="CASCADE"), nullable=False, unique=True)
    description = Column(Text)
    shot_type = Column(ENUM("close_up", "medium_shot", "wide_shot", "extreme_close_up", "establishing_shot", "aerial", "tracking", "static", "pan", "tilt", "drone", "360", name="shot_type_enum"))
    camera_angle = Column(ENUM("front", "rear", "side_left", "side_right", "top_down", "low_angle", "eye_level", "three_quarter", "interior", "detail", name="camera_angle_enum"))
    subject_part = Column(Text)
    mood = Column(ENUM("sporty", "elegant", "technical", "luxury", "adventure", "minimal", "dynamic", "serene", name="mood_enum"))
    dominant_colours = Column(ARRAY(Text))
    confidence_score = Column(Float)
    review_flag = Column(Boolean, default=False)
    embedding_768 = Column(Vector(768) if _TRY_PGVECTOR else JSONB)
    embedding_model = Column(Text, default="openai/clip-vit-large-patch14")
    extracted_at = Column(DateTime(timezone=True), server_default=func.now())
    asset = relationship("Asset", back_populates="metadata")


class VideoEvent(Base, TimestampMixin):
    __tablename__ = "video_events"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    asset_id = Column(UUID(as_uuid=True), ForeignKey("assets.id", ondelete="CASCADE"), nullable=False)
    start_ms = Column(Integer, nullable=False)
    end_ms = Column(Integer, nullable=False)
    duration_ms = Column(Integer)
    description = Column(Text)
    shot_type = Column(ENUM("close_up", "medium_shot", "wide_shot", "extreme_close_up", "establishing_shot", "aerial", "tracking", "static", "pan", "tilt", "drone", "360", name="shot_type_enum"))
    camera_angle = Column(ENUM("front", "rear", "side_left", "side_right", "top_down", "low_angle", "eye_level", "three_quarter", "interior", "detail", name="camera_angle_enum"))
    subject_part = Column(Text)
    mood = Column(ENUM("sporty", "elegant", "technical", "luxury", "adventure", "minimal", "dynamic", "serene", name="mood_enum"))
    embedding_768 = Column(Vector(768) if _TRY_PGVECTOR else JSONB)
    confidence_score = Column(Float)
    keyframe_path = Column(Text)
    asset = relationship("Asset", back_populates="video_events")


class BrandConfig(Base, TimestampMixin):
    __tablename__ = "brand_configs"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(Text, nullable=False)
    version = Column(Integer, nullable=False, default=1)
    colours = Column(JSONB, nullable=False, default=dict)
    typography = Column(JSONB, nullable=False, default=dict)
    tone_of_voice = Column(JSONB, nullable=False, default=dict)
    approved_terminology = Column(ARRAY(Text))
    restricted_visuals = Column(ARRAY(Text))
    logo_paths = Column(JSONB)


class BrochureNode(Base, TimestampMixin):
    __tablename__ = "brochure_nodes"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    section = Column(ENUM("Exterior", "Interior", "Performance", "Safety", "Technology", "Comfort", "Sustainability", "Brand_Story", name="section_enum"), nullable=False)
    title = Column(Text, nullable=False)
    content = Column(Text, nullable=False)
    key_features = Column(ARRAY(Text))
    taglines = Column(ARRAY(Text))
    spec_highlights = Column(JSONB)
    car_part_referenced = Column(ARRAY(Text))
    tone_tags = Column(ARRAY(Text))
    embedding_768 = Column(Vector(768) if _TRY_PGVECTOR else JSONB)
    embedding_model = Column(Text, default="openai/clip-vit-large-patch14")
    page_number = Column(Integer)
    source_pdf = Column(Text)
    mappings = relationship("BrochureAssetMap", back_populates="brochure_node")
    captions = relationship("CaptionLibrary", back_populates="brochure_node")
    voiceovers = relationship("VoiceoverLibrary", back_populates="brochure_node")


class BrochureAssetMap(Base, TimestampMixin):
    __tablename__ = "brochure_asset_map"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    brochure_node_id = Column(UUID(as_uuid=True), ForeignKey("brochure_nodes.id", ondelete="CASCADE"), nullable=False)
    asset_id = Column(UUID(as_uuid=True), ForeignKey("assets.id", ondelete="CASCADE"), nullable=False)
    video_event_id = Column(UUID(as_uuid=True), ForeignKey("video_events.id", ondelete="SET NULL"))
    similarity_score = Column(Float, nullable=False)
    mapping_type = Column(ENUM("semantic", "rule_based", "manual", "override", name="mapping_type_enum"), default="semantic")
    confidence_score = Column(Float, nullable=False)
    is_approved = Column(Boolean, default=None)
    reviewer_notes = Column(Text)
    rank = Column(Integer, nullable=False)
    brochure_node = relationship("BrochureNode", back_populates="mappings")


class CaptionLibrary(Base, TimestampMixin):
    __tablename__ = "captions_library"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    brochure_node_id = Column(UUID(as_uuid=True), ForeignKey("brochure_nodes.id", ondelete="SET NULL"))
    car_part = Column(Text)
    tone = Column(ENUM("sporty", "elegant", "technical", "luxury", "adventure", "minimal", "dynamic", "serene", name="tone_enum"), nullable=False)
    duration_class = Column(ENUM("short", "medium", "long", name="duration_class_enum"), nullable=False, default="medium")
    text = Column(Text, nullable=False)
    word_count = Column(Integer)
    is_brand_compliant = Column(Boolean, default=True)
    compliance_notes = Column(Text)
    usage_count = Column(Integer, default=0)
    brochure_node = relationship("BrochureNode", back_populates="captions")


class VoiceoverLibrary(Base, TimestampMixin):
    __tablename__ = "voiceover_library"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    brochure_node_id = Column(UUID(as_uuid=True), ForeignKey("brochure_nodes.id", ondelete="SET NULL"))
    car_part = Column(Text)
    tone = Column(ENUM("sporty", "elegant", "technical", "luxury", "adventure", "minimal", "dynamic", "serene", name="tone_enum"), nullable=False)
    duration_class = Column(ENUM("short", "medium", "long", name="duration_class_enum"), nullable=False, default="medium")
    text = Column(Text, nullable=False)
    estimated_duration_ms = Column(Integer)
    is_brand_compliant = Column(Boolean, default=True)
    compliance_notes = Column(Text)
    usage_count = Column(Integer, default=0)
    brochure_node = relationship("BrochureNode", back_populates="voiceovers")


class ReelRequest(Base, TimestampMixin):
    __tablename__ = "reel_requests"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    user_query = Column(Text, nullable=False)
    duration_target = Column(ENUM("10s", "20s", "30s", "60s", "custom", name="duration_target_enum"), nullable=False)
    duration_ms = Column(Integer)
    platform = Column(ENUM("instagram_reels", "tiktok", "youtube_shorts", "linkedin", "twitter", "facebook", "custom", name="platform_enum"), nullable=False)
    tone = Column(ENUM("sporty", "elegant", "technical", "luxury", "adventure", "minimal", "dynamic", "serene", name="tone_enum"), nullable=False)
    aspect_ratio = Column(ENUM("9:16", "16:9", "1:1", "4:5", name="aspect_ratio_enum"), nullable=False, default="9:16")
    brand_config_id = Column(UUID(as_uuid=True), ForeignKey("brand_configs.id"))
    additional_constraints = Column(JSONB)
    status = Column(ENUM("pending", "planning", "generating", "review", "completed", "failed", name="reel_status_enum"), nullable=False, default="pending")
    error_message = Column(Text)
    script = relationship("ReelScript", uselist=False, back_populates="reel_request")
    manifest = relationship("ReelManifest", uselist=False, back_populates="reel_request")


class ReelScript(Base, TimestampMixin):
    __tablename__ = "reel_scripts"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    reel_request_id = Column(UUID(as_uuid=True), ForeignKey("reel_requests.id", ondelete="CASCADE"), nullable=False, unique=True)
    script_json = Column(JSONB, nullable=False)
    total_beats = Column(Integer, nullable=False)
    total_duration_ms = Column(Integer, nullable=False)
    validation_status = Column(ENUM("pending", "valid", "invalid", "corrected", name="validation_status_enum"), nullable=False, default="pending")
    validation_errors = Column(ARRAY(Text))
    generation_attempts = Column(Integer, default=1)
    reel_request = relationship("ReelRequest", back_populates="script")


class BeatAssetCandidate(Base, TimestampMixin):
    __tablename__ = "beat_asset_candidates"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    reel_request_id = Column(UUID(as_uuid=True), ForeignKey("reel_requests.id", ondelete="CASCADE"), nullable=False)
    beat_number = Column(Integer, nullable=False)
    beat_intent = Column(Text, nullable=False)
    car_parts = Column(ARRAY(Text))
    tone = Column(ENUM("sporty", "elegant", "technical", "luxury", "adventure", "minimal", "dynamic", "serene", name="tone_enum"))
    asset_id = Column(UUID(as_uuid=True), ForeignKey("assets.id", ondelete="SET NULL"))
    video_event_id = Column(UUID(as_uuid=True), ForeignKey("video_events.id", ondelete="SET NULL"))
    similarity_score = Column(Float)
    rank = Column(Integer, nullable=False)
    mapping_source = Column(ENUM("semantic", "brochure_map", "rule_based", "fallback", name="mapping_source_enum"), default="semantic")
    is_selected = Column(Boolean, default=False)


class ReelManifest(Base, TimestampMixin):
    __tablename__ = "reel_manifests"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    reel_request_id = Column(UUID(as_uuid=True), ForeignKey("reel_requests.id", ondelete="CASCADE"), nullable=False, unique=True)
    manifest_json = Column(JSONB, nullable=False)
    total_duration_ms = Column(Integer, nullable=False)
    beat_count = Column(Integer, nullable=False)
    asset_count = Column(Integer, nullable=False)
    is_validated = Column(Boolean, default=False)
    validation_report = Column(JSONB)
    render_status = Column(ENUM("pending", "rendering", "completed", "failed", name="render_status_enum"), nullable=False, default="pending")
    render_output_path = Column(Text)
    remotion_export_path = Column(Text)
    reel_request = relationship("ReelRequest", back_populates="manifest")


class AuditLog(Base):
    __tablename__ = "audit_log"
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    table_name = Column(Text, nullable=False)
    record_id = Column(UUID(as_uuid=True), nullable=False)
    action = Column(ENUM("INSERT", "UPDATE", "DELETE", name="audit_action_enum"), nullable=False)
    old_data = Column(JSONB)
    new_data = Column(JSONB)
    performed_by = Column(Text)
    performed_at = Column(DateTime(timezone=True), server_default=func.now())


def get_engine(database_url: Optional[str] = None):
    url = database_url or os.environ.get("DATABASE_URL", "postgresql://reel_user:reel_password@localhost:5432/reel_creator")
    return create_engine(url, future=True)


def init_db(engine=None):
    engine = engine or get_engine()
    Base.metadata.create_all(engine)
    print("Database tables created / verified.")


if __name__ == "__main__":
    init_db()