backend/db/models.py

"""SQLAlchemy ORM models for AutoForge."""

import uuid
from datetime import datetime, timezone

from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text
from sqlalchemy.orm import relationship

from backend.db.database import Base


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _generate_uuid() -> str:
    return str(uuid.uuid4())


class User(Base):
    """User account."""

    __tablename__ = "users"

    id = Column(String(36), primary_key=True, default=_generate_uuid)
    username = Column(String(80), unique=True, nullable=False, index=True)
    email = Column(String(255), unique=True, nullable=False, index=True)
    hashed_password = Column(String(255), nullable=False)
    created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False)

    workflows = relationship("Workflow", back_populates="user", lazy="selectin")

    def __repr__(self) -> str:
        return f"<User {self.username}>"


class Workflow(Base):
    """A code-generation workflow initiated by a user."""

    __tablename__ = "workflows"

    id = Column(String(36), primary_key=True, default=_generate_uuid)
    user_id = Column(String(36), ForeignKey("users.id"), nullable=False, index=True)
    prompt = Column(Text, nullable=False)
    status = Column(String(20), nullable=False, default="pending")
    created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False)
    updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow, nullable=False)

    user = relationship("User", back_populates="workflows")
    tasks = relationship("Task", back_populates="workflow", lazy="selectin")
    events = relationship("Event", back_populates="workflow", lazy="selectin")

    def __repr__(self) -> str:
        return f"<Workflow {self.id} [{self.status}]>"


class Task(Base):
    """A single task within a workflow DAG."""

    __tablename__ = "tasks"

    id = Column(String(36), primary_key=True, default=_generate_uuid)
    workflow_id = Column(String(36), ForeignKey("workflows.id"), nullable=False, index=True)
    task_type = Column(String(50), nullable=False)
    status = Column(String(20), nullable=False, default="pending")
    input_data = Column(Text, nullable=True)
    output_data = Column(Text, nullable=True)
    created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False)

    workflow = relationship("Workflow", back_populates="tasks")
    attempts = relationship("Attempt", back_populates="task", lazy="selectin")

    def __repr__(self) -> str:
        return f"<Task {self.id} [{self.task_type}]>"


class Attempt(Base):
    """A single execution attempt for a task (supports retry logic)."""

    __tablename__ = "attempts"

    id = Column(String(36), primary_key=True, default=_generate_uuid)
    task_id = Column(String(36), ForeignKey("tasks.id"), nullable=False, index=True)
    attempt_number = Column(Integer, nullable=False, default=1)
    status = Column(String(20), nullable=False, default="pending")
    error_message = Column(Text, nullable=True)
    created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False)

    task = relationship("Task", back_populates="attempts")

    def __repr__(self) -> str:
        return f"<Attempt {self.id} #{self.attempt_number}>"


class Event(Base):
    """A workflow event for streaming to the frontend."""

    __tablename__ = "events"

    id = Column(String(36), primary_key=True, default=_generate_uuid)
    workflow_id = Column(String(36), ForeignKey("workflows.id"), nullable=False, index=True)
    event_type = Column(String(50), nullable=False)
    payload = Column(Text, nullable=True)
    created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False)

    workflow = relationship("Workflow", back_populates="events")

    def __repr__(self) -> str:
        return f"<Event {self.id} [{self.event_type}]>"