spoolman2/spoolman/database/print_job.py

"""Database operations for print jobs."""

from datetime import datetime
from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import joinedload

from spoolman.database.models import PrintJob, Spool, SpoolAdjustment
from spoolman.exceptions import ItemNotFoundError


async def create(
    *,
    db: AsyncSession,
    spool_id: int,
    filename: str,
    filament_used_g: float,
    filament_used_mm: Optional[float] = None,
) -> PrintJob:
    """Create a new pending print job."""
    # Verify spool exists
    spool = await db.get(Spool, spool_id)
    if spool is None:
        raise ItemNotFoundError(f"Spool with ID {spool_id} not found.")

    job = PrintJob(
        spool_id=spool_id,
        created=datetime.utcnow(),
        filename=filename,
        filament_used_g=filament_used_g,
        filament_used_mm=filament_used_mm,
        status="pending",
    )
    db.add(job)
    await db.commit()
    await db.refresh(job)

    # Load spool relationship
    result = await db.execute(
        select(PrintJob).where(PrintJob.id == job.id).options(joinedload(PrintJob.spool))
    )
    return result.scalar_one()


async def get_by_id(db: AsyncSession, job_id: int) -> PrintJob:
    """Get a print job by ID."""
    result = await db.execute(
        select(PrintJob).where(PrintJob.id == job_id).options(joinedload(PrintJob.spool))
    )
    job = result.scalar_one_or_none()
    if job is None:
        raise ItemNotFoundError(f"Print job with ID {job_id} not found.")
    return job


async def find(
    *,
    db: AsyncSession,
    status: Optional[str] = None,
    spool_id: Optional[int] = None,
    limit: Optional[int] = None,
    offset: int = 0,
) -> tuple[list[PrintJob], int]:
    """Find print jobs with optional filters."""
    query = select(PrintJob).options(joinedload(PrintJob.spool))

    if status:
        query = query.where(PrintJob.status == status)
    if spool_id:
        query = query.where(PrintJob.spool_id == spool_id)

    # Get total count
    count_query = select(PrintJob)
    if status:
        count_query = count_query.where(PrintJob.status == status)
    if spool_id:
        count_query = count_query.where(PrintJob.spool_id == spool_id)
    count_result = await db.execute(count_query)
    total = len(count_result.all())

    # Apply ordering and pagination
    query = query.order_by(PrintJob.created.desc())
    if offset:
        query = query.offset(offset)
    if limit:
        query = query.limit(limit)

    result = await db.execute(query)
    return list(result.scalars().all()), total


async def complete(db: AsyncSession, job_id: int, comment: Optional[str] = None) -> PrintJob:
    """Complete a print job - deduct filament from spool."""
    job = await get_by_id(db, job_id)
    if job.status != "pending":
        raise ValueError(f"Job {job_id} is not pending (status: {job.status}).")

    # Deduct filament from spool
    spool = await db.get(Spool, job.spool_id)
    if spool is None:
        raise ItemNotFoundError(f"Spool with ID {job.spool_id} not found.")

    spool.used_weight += job.filament_used_g
    if spool.first_used is None:
        spool.first_used = datetime.utcnow()
    spool.last_used = datetime.utcnow()

    # Record adjustment
    adjustment = SpoolAdjustment(
        spool_id=spool.id,
        timestamp=datetime.utcnow(),
        adjustment_type="weight",
        value=-job.filament_used_g,
        comment=comment or f"Print job: {job.filename}",
    )
    db.add(adjustment)

    # Update job status
    job.status = "completed"
    job.finished = datetime.utcnow()

    await db.commit()

    # Reload with relationships
    return await get_by_id(db, job_id)


async def cancel(db: AsyncSession, job_id: int) -> PrintJob:
    """Cancel a print job - flag spool for weighing."""
    job = await get_by_id(db, job_id)
    if job.status != "pending":
        raise ValueError(f"Job {job_id} is not pending (status: {job.status}).")

    # Flag spool for weighing
    spool = await db.get(Spool, job.spool_id)
    if spool is None:
        raise ItemNotFoundError(f"Spool with ID {job.spool_id} not found.")

    spool.needs_weighing = True

    # Update job status
    job.status = "cancelled"
    job.finished = datetime.utcnow()

    await db.commit()

    # Reload with relationships
    return await get_by_id(db, job_id)


async def delete(db: AsyncSession, job_id: int) -> None:
    """Delete a print job."""
    job = await get_by_id(db, job_id)
    await db.delete(job)
    await db.commit()