"""YOLOv8-based person detector for peepfilt."""

from __future__ import annotations
from typing import Iterable, List

import torch
from ultralytics import YOLO

from peepfilt.config import (
    MODEL_NAME,
    PERSON_CLASS_ID,
    IMAGE_SIZE,
    DEFAULT_CONF,
    DEFAULT_BATCH_SIZE,
)
from peepfilt.utils import get_logger

logger = get_logger("peepfilt.detector")


def get_default_device() -> str:
    """Return 'cuda' if available, otherwise 'cpu'."""
    return "cuda" if torch.cuda.is_available() else "cpu"


class PersonDetector:
    """YOLOv8-based detector for person presence in images."""

    def __init__(self, model_name: str = MODEL_NAME, device: str | None = None):
        self.device = device or get_default_device()
        logger.info(f"Loading model {model_name} on device={self.device}")
        self.model = YOLO(model_name)
        self.model.to(self.device)

    def predict_paths(
        self,
        image_paths: Iterable[str],
        conf: float = DEFAULT_CONF,
        batch_size: int = DEFAULT_BATCH_SIZE,
    ) -> List[bool]:
        """Return list of booleans: True if person detected, False otherwise."""
        results = self.model.predict(
            source=list(image_paths),
            conf=conf,
            imgsz=IMAGE_SIZE,
            classes=[PERSON_CLASS_ID],
            verbose=False,
            batch=batch_size,
            device=self.device,
        )
        return [len(r.boxes) > 0 for r in results]
