modcomms/backend/app/models/models.py

import uuid
from datetime import datetime
from typing import Optional

from sqlalchemy import Boolean, DateTime, ForeignKey, Integer, String, Text, UniqueConstraint, func
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import Mapped, mapped_column, relationship

from app.models.database import Base


class Agency(Base):
    """Agency/organization that users belong to."""
    __tablename__ = "agencies"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    users: Mapped[list["User"]] = relationship("User", back_populates="agency")
    campaigns: Mapped[list["Campaign"]] = relationship("Campaign", back_populates="agency")


class User(Base):
    """User account linked to Azure AD."""
    __tablename__ = "users"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    azure_ad_oid: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    email: Mapped[str] = mapped_column(String(255), nullable=False)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    role: Mapped[str] = mapped_column(String(50), nullable=False, default="basic_user")
    agency_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("agencies.id"), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    agency: Mapped[Optional["Agency"]] = relationship("Agency", back_populates="users")
    campaigns: Mapped[list["Campaign"]] = relationship("Campaign", back_populates="created_by_user")
    proofs: Mapped[list["Proof"]] = relationship("Proof", back_populates="created_by_user")
    flagged_items: Mapped[list["FlaggedItem"]] = relationship("FlaggedItem", back_populates="submitter")
    resolved_items: Mapped[list["ResolvedItem"]] = relationship("ResolvedItem", back_populates="submitter")


class Campaign(Base):
    """Marketing campaign containing proofs."""
    __tablename__ = "campaigns"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    workfront_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    client_lead: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    agency_lead: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    brand_guidelines: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
    status: Mapped[str] = mapped_column(String(50), default="In Progress")
    agency_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("agencies.id"), nullable=True)
    created_by: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())

    # Relationships
    agency: Mapped[Optional["Agency"]] = relationship("Agency", back_populates="campaigns")
    created_by_user: Mapped[Optional["User"]] = relationship("User", back_populates="campaigns")
    proofs: Mapped[list["Proof"]] = relationship("Proof", back_populates="campaign", cascade="all, delete-orphan")


class Proof(Base):
    """Marketing proof/asset to be reviewed."""
    __tablename__ = "proofs"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    campaign_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("campaigns.id", ondelete="CASCADE"), nullable=False)
    proof_name: Mapped[str] = mapped_column(String(255), nullable=False)
    channel: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    sub_channel: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    proof_type: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    workfront_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    created_by: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    campaign: Mapped["Campaign"] = relationship("Campaign", back_populates="proofs")
    created_by_user: Mapped[Optional["User"]] = relationship("User", back_populates="proofs")
    versions: Mapped[list["ProofVersion"]] = relationship("ProofVersion", back_populates="proof", cascade="all, delete-orphan", order_by="desc(ProofVersion.version)")

    __table_args__ = (
        UniqueConstraint("campaign_id", "proof_name", name="uq_campaign_proof_name"),
    )


class ProofVersion(Base):
    """Version of a proof with analysis results."""
    __tablename__ = "proof_versions"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    proof_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proofs.id", ondelete="CASCADE"), nullable=False)
    version: Mapped[int] = mapped_column(Integer, nullable=False)
    file_storage_key: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
    thumbnail_url: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    agent_review: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
    overall_status: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
    workfront_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    file_hash: Mapped[Optional[str]] = mapped_column(String(32), nullable=True)
    is_identical_file: Mapped[Optional[bool]] = mapped_column(Boolean, nullable=True, default=False)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    proof: Mapped["Proof"] = relationship("Proof", back_populates="versions")
    flagged_items: Mapped[list["FlaggedItem"]] = relationship("FlaggedItem", back_populates="proof_version")
    resolved_items: Mapped[list["ResolvedItem"]] = relationship("ResolvedItem", back_populates="proof_version")
    error_items: Mapped[list["ErrorItem"]] = relationship("ErrorItem", back_populates="proof_version")

    __table_args__ = (
        UniqueConstraint("proof_id", "version", name="uq_proof_version"),
    )


class FlaggedItem(Base):
    """Record of a flagged issue on a proof version."""
    __tablename__ = "flagged_items"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    proof_version_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proof_versions.id"), nullable=False)
    agent_flagged: Mapped[str] = mapped_column(String(100), nullable=False)
    comments: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    submitter_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    proof_version: Mapped["ProofVersion"] = relationship("ProofVersion", back_populates="flagged_items")
    submitter: Mapped[Optional["User"]] = relationship("User", back_populates="flagged_items")


class ResolvedItem(Base):
    """Record of a resolved issue on a proof version."""
    __tablename__ = "resolved_items"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    proof_version_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proof_versions.id"), nullable=False)
    agent: Mapped[str] = mapped_column(String(100), nullable=False)
    issue: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    resolution: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    submitter_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    proof_version: Mapped["ProofVersion"] = relationship("ProofVersion", back_populates="resolved_items")
    submitter: Mapped[Optional["User"]] = relationship("User", back_populates="resolved_items")


class ErrorItem(Base):
    """Record of an analysis error on a proof version."""
    __tablename__ = "error_items"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    proof_version_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proof_versions.id"), nullable=False)
    error_summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    proof_version: Mapped["ProofVersion"] = relationship("ProofVersion", back_populates="error_items")


class DropdownOption(Base):
    """Configurable dropdown options for channels/sub-channels/proof types."""
    __tablename__ = "dropdown_options"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    option_type: Mapped[str] = mapped_column(String(50), nullable=False)  # 'channel', 'sub_channel', 'proof_type'
    parent_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("dropdown_options.id"), nullable=True)
    value: Mapped[str] = mapped_column(String(255), nullable=False)
    display_order: Mapped[int] = mapped_column(Integer, default=0)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Self-referential relationship for hierarchy
    parent: Mapped[Optional["DropdownOption"]] = relationship("DropdownOption", remote_side=[id], back_populates="children")
    children: Mapped[list["DropdownOption"]] = relationship(
        "DropdownOption",
        back_populates="parent",
        order_by="DropdownOption.display_order"
    )


class KnowledgeBase(Base):
    """An agent knowledge base entry (one per agent spec type)."""
    __tablename__ = "knowledge_bases"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    agent_key: Mapped[str] = mapped_column(String(100), unique=True, nullable=False)
    display_name: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    source_documents: Mapped[list["SourceDocument"]] = relationship(
        "SourceDocument", back_populates="knowledge_base", cascade="all, delete-orphan"
    )
    spec_versions: Mapped[list["SpecVersion"]] = relationship(
        "SpecVersion", back_populates="knowledge_base", cascade="all, delete-orphan",
        order_by="desc(SpecVersion.version_number)"
    )
    processing_jobs: Mapped[list["ProcessingJob"]] = relationship(
        "ProcessingJob", back_populates="knowledge_base", cascade="all, delete-orphan",
        order_by="desc(ProcessingJob.created_at)"
    )


class SourceDocument(Base):
    """An uploaded source document linked to a knowledge base."""
    __tablename__ = "source_documents"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False
    )
    filename: Mapped[str] = mapped_column(String(500), nullable=False)
    file_storage_key: Mapped[str] = mapped_column(String(500), nullable=False)
    file_size_bytes: Mapped[int] = mapped_column(Integer, nullable=False)
    mime_type: Mapped[str] = mapped_column(String(255), nullable=False)
    uploaded_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    uploaded_by_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    parsed_markdown: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    parse_status: Mapped[str] = mapped_column(String(50), default="pending", nullable=False)
    parse_error: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    knowledge_base: Mapped["KnowledgeBase"] = relationship("KnowledgeBase", back_populates="source_documents")


class SpecVersion(Base):
    """A generated spec document with version history."""
    __tablename__ = "spec_versions"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False
    )
    version_number: Mapped[int] = mapped_column(Integer, nullable=False)
    content: Mapped[str] = mapped_column(Text, nullable=False)
    source_document_ids: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
    generated_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    generated_by_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    processing_job_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True), ForeignKey("processing_jobs.id"), nullable=True
    )
    is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
    char_count: Mapped[int] = mapped_column(Integer, nullable=False)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    knowledge_base: Mapped["KnowledgeBase"] = relationship("KnowledgeBase", back_populates="spec_versions")
    processing_job: Mapped[Optional["ProcessingJob"]] = relationship(
        "ProcessingJob", back_populates="spec_version",
        foreign_keys="[SpecVersion.processing_job_id]",
    )

    __table_args__ = (
        UniqueConstraint("knowledge_base_id", "version_number", name="uq_kb_version_number"),
    )


class ProcessingJob(Base):
    """Tracks a document processing pipeline run."""
    __tablename__ = "processing_jobs"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False
    )
    status: Mapped[str] = mapped_column(String(50), default="pending", nullable=False)
    triggered_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    triggered_by_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    total_documents: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    parsed_documents: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    spec_version_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("spec_versions.id"), nullable=True)
    error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    log: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
    started_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
    completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    # Relationships
    knowledge_base: Mapped["KnowledgeBase"] = relationship("KnowledgeBase", back_populates="processing_jobs")
    spec_version: Mapped[Optional["SpecVersion"]] = relationship(
        "SpecVersion", back_populates="processing_job",
        foreign_keys="[ProcessingJob.spec_version_id]",
    )