spoolman2/spoolman/database/location.py

"""Database operations for Location model."""

from typing import Optional

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from spoolman.database import models
from spoolman.exceptions import ItemNotFoundError


async def create(
    *,
    db: AsyncSession,
    name: str,
    parent_id: Optional[int] = None,
    location_type: Optional[str] = None,
    description: Optional[str] = None,
) -> models.Location:
    """Create a new location."""
    # Verify parent exists if specified
    if parent_id is not None:
        parent = await db.get(models.Location, parent_id)
        if parent is None:
            raise ItemNotFoundError(f"Parent location with ID {parent_id} not found.")

    db_item = models.Location(
        name=name,
        parent_id=parent_id,
        location_type=location_type,
        description=description,
    )
    db.add(db_item)
    await db.commit()
    return db_item


async def get_by_id(db: AsyncSession, location_id: int) -> models.Location:
    """Get a location by ID."""
    stmt = select(models.Location).where(models.Location.id == location_id)
    result = await db.execute(stmt)
    db_item = result.scalar_one_or_none()
    if db_item is None:
        raise ItemNotFoundError(f"Location with ID {location_id} not found.")
    return db_item


async def find(
    *,
    db: AsyncSession,
    parent_id: Optional[int] = None,
) -> list[models.Location]:
    """Find locations, optionally filtered by parent."""
    stmt = select(models.Location).order_by(models.Location.name)
    if parent_id is not None:
        stmt = stmt.where(models.Location.parent_id == parent_id)
    result = await db.execute(stmt)
    return list(result.scalars().all())


async def find_all(db: AsyncSession) -> list[models.Location]:
    """Get all locations."""
    stmt = select(models.Location).order_by(models.Location.name)
    result = await db.execute(stmt)
    return list(result.scalars().all())


async def update(
    *,
    db: AsyncSession,
    location_id: int,
    data: dict,
) -> models.Location:
    """Update a location."""
    db_item = await get_by_id(db, location_id)

    # Verify parent exists if changing parent
    if "parent_id" in data and data["parent_id"] is not None:
        # Prevent setting parent to self
        if data["parent_id"] == location_id:
            raise ValueError("Location cannot be its own parent.")
        parent = await db.get(models.Location, data["parent_id"])
        if parent is None:
            raise ItemNotFoundError(f"Parent location with ID {data['parent_id']} not found.")

    for key, value in data.items():
        setattr(db_item, key, value)

    await db.commit()
    return db_item


async def delete(db: AsyncSession, location_id: int) -> None:
    """Delete a location."""
    db_item = await get_by_id(db, location_id)
    await db.delete(db_item)
    await db.commit()


async def get_full_path(db: AsyncSession, location_id: int) -> str:
    """Get the full path of a location as a string (e.g., 'Room > Shelf > Bin')."""
    parts = []
    current_id: Optional[int] = location_id

    while current_id is not None:
        location = await get_by_id(db, current_id)
        parts.insert(0, location.name)
        current_id = location.parent_id

    return " > ".join(parts)


async def get_tree(db: AsyncSession) -> list[dict]:
    """Get all locations as a nested tree structure."""
    all_locations = await find_all(db)

    # Build lookup dict
    by_id = {loc.id: loc for loc in all_locations}

    # Build tree
    def build_node(loc: models.Location) -> dict:
        return {
            "id": loc.id,
            "name": loc.name,
            "location_type": loc.location_type,
            "description": loc.description,
            "children": [
                build_node(by_id[child.id])
                for child in all_locations
                if child.parent_id == loc.id
            ],
        }

    # Get root nodes (no parent)
    roots = [loc for loc in all_locations if loc.parent_id is None]
    return [build_node(loc) for loc in roots]


async def get_spool_count(db: AsyncSession, location_id: int) -> int:
    """Get the number of spools at this location (not including children)."""
    stmt = (
        select(func.count(models.Spool.id))
        .where(models.Spool.location_id == location_id)
        .where((models.Spool.archived == False) | (models.Spool.archived == None))  # noqa: E712
    )
    result = await db.execute(stmt)
    return result.scalar() or 0