from typing import List, Optional, Union

from sqlalchemy.orm import Session, Query, exc, load_only

from watcher.db import Staff
from watcher.logic.exceptions import StaffNotFound
from watcher.crud.base import get_object_by_model


def get_staff_by_uid(db: Session, uid: Union[int, str]) -> Staff:
    try:
        return (
            db.query(Staff)
            .filter(Staff.uid == str(uid))
            .options(load_only(Staff.login, Staff.user_id, Staff.uid))
        ).one()
    except exc.NoResultFound:
        raise StaffNotFound


def get_staff_by_id(db: Session, staff_id: int) -> Staff:
    return get_object_by_model(db=db, model=Staff, object_id=staff_id)


def query_staff_by_ids(db: Session, staff_ids: List[int], query: Optional[Query] = None) -> Query:
    if not query:
        query = db.query(Staff)

    return query.filter(Staff.id.in_(staff_ids))
