models.py

"""SQLAlchemy models for the Annotation Management System."""

import json
from datetime import datetime, timezone
from flask_sqlalchemy import SQLAlchemy

db = SQLAlchemy()


class User(db.Model):
    __tablename__ = 'users'

    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    role = db.Column(db.String(20), nullable=False, default='annotator')  # 'admin' or 'annotator'
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))

    assignments = db.relationship('Assignment', backref='user', lazy=True, cascade='all, delete-orphan')

    def to_dict(self):
        return {
            'id': self.id,
            'username': self.username,
            'role': self.role,
            'created_at': self.created_at.isoformat() if self.created_at else None,
        }


class CsvFile(db.Model):
    __tablename__ = 'csv_files'

    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(255), unique=True, nullable=False)
    filepath = db.Column(db.String(512), nullable=False)
    total_samples = db.Column(db.Integer, default=0)
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))

    assignments = db.relationship('Assignment', backref='csv_file', lazy=True, cascade='all, delete-orphan')

    def to_dict(self):
        return {
            'id': self.id,
            'filename': self.filename,
            'total_samples': self.total_samples,
            'created_at': self.created_at.isoformat() if self.created_at else None,
        }


class Assignment(db.Model):
    __tablename__ = 'assignments'

    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    file_id = db.Column(db.Integer, db.ForeignKey('csv_files.id'), nullable=False)
    assigned_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))

    reviews = db.relationship('Review', backref='assignment', lazy=True, cascade='all, delete-orphan')
    annotation_history = db.relationship(
        'AnnotationHistory', backref='assignment', lazy=True, cascade='all, delete-orphan'
    )

    # Prevent duplicate assignments
    __table_args__ = (db.UniqueConstraint('user_id', 'file_id', name='uq_user_file'),)

    def to_dict(self):
        reviewed_count = Review.query.filter_by(assignment_id=self.id).count()
        total = self.csv_file.total_samples if self.csv_file else 0
        return {
            'id': self.id,
            'user_id': self.user_id,
            'username': self.user.username if self.user else None,
            'file_id': self.file_id,
            'filename': self.csv_file.filename if self.csv_file else None,
            'total_samples': total,
            'reviewed_count': reviewed_count,
            'progress': round((reviewed_count / total * 100), 1) if total > 0 else 0,
            'assigned_at': self.assigned_at.isoformat() if self.assigned_at else None,
        }


class Review(db.Model):
    __tablename__ = 'reviews'

    id = db.Column(db.Integer, primary_key=True)
    assignment_id = db.Column(db.Integer, db.ForeignKey('assignments.id'), nullable=False)
    sample_index = db.Column(db.Integer, nullable=False)
    status = db.Column(db.String(20), nullable=False)  # 'accept' or 'reject'
    comment = db.Column(db.Text, default='')
    # JSON string storing annotator-edited annotation data (emotion/distortions/empathy)
    annotations_edit = db.Column(db.Text, default=None)
    reviewed_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))

    # One review per sample per assignment
    __table_args__ = (db.UniqueConstraint('assignment_id', 'sample_index', name='uq_assignment_sample'),)

    def to_dict(self):
        return {
            'id': self.id,
            'assignment_id': self.assignment_id,
            'sample_index': self.sample_index,
            'status': self.status,
            'comment': self.comment,
            'annotations_edit': self.annotations_edit,  # raw JSON string
            'reviewed_at': self.reviewed_at.isoformat() if self.reviewed_at else None,
        }


class AnnotationHistory(db.Model):
    """Immutable snapshots created whenever an annotation is changed."""

    __tablename__ = 'annotation_history'

    id = db.Column(db.Integer, primary_key=True)
    assignment_id = db.Column(db.Integer, db.ForeignKey('assignments.id'), nullable=False)
    sample_index = db.Column(db.Integer, nullable=False)
    annotations_json = db.Column(db.Text, nullable=False)
    edited_by = db.Column(db.String(80), nullable=False)
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc), nullable=False)

    __table_args__ = (
        db.Index(
            'ix_annotation_history_lookup',
            'assignment_id', 'sample_index', 'created_at',
        ),
    )

    def to_dict(self):
        try:
            annotations = json.loads(self.annotations_json)
        except (TypeError, json.JSONDecodeError):
            annotations = None
        created_at = self.created_at
        if created_at and created_at.tzinfo is None:
            created_at = created_at.replace(tzinfo=timezone.utc)
        return {
            'id': self.id,
            'assignment_id': self.assignment_id,
            'sample_index': self.sample_index,
            'annotations': annotations,
            'edited_by': self.edited_by,
            'created_at': created_at.isoformat() if created_at else None,
        }