"""Directory batch processing utilities for peepfilt."""

from __future__ import annotations
import shutil
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

from peepfilt.config import (
    IMAGE_EXTENSIONS,
    DEFAULT_CONF,
    DEFAULT_BATCH_SIZE,
)
from peepfilt.utils import get_logger

logger = get_logger("peepfilt.batch")


def _iter_images_once(directory: Path) -> List[Path]:
    """Return list of image files in a directory (non-recursive)."""
    return [
        p
        for p in sorted(directory.iterdir())
        if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS
    ]


def _iter_images_recursive(root_dir: Path) -> List[Tuple[Path, List[Path]]]:
    """Yield (directory, images_list) for each directory under root_dir.

    Skips any directory named 'no_people'.
    """
    results = []
    for dirpath in sorted(root_dir.rglob("*")):
        if not dirpath.is_dir():
            continue
        if dirpath.name == "no_people":
            continue
        images = _iter_images_once(dirpath)
        if images:
            results.append((dirpath, images))
    return results


def _chunks(seq: List[Path], size: int) -> Iterable[List[Path]]:
    """Yield successive chunks of `seq` with max length `size`."""
    for i in range(0, len(seq), size):
        yield seq[i : i + size]


def scan_and_filter(
    input_dir: Path | str,
    no_people_dir: Path | str | None = None,
    conf: float = DEFAULT_CONF,
    batch_size: int = DEFAULT_BATCH_SIZE,
    recursive: bool = True,
) -> Dict[int | str, str]:
    """Scan images in a directory; move those without people to no_people_dir.

    If recursive=True: scans all subdirectories (excluding 'no_people' dirs).
    If no_people_dir is specified: moves all no-people images to that fixed directory.
    If no_people_dir is not specified: creates 'no_people' subdirectories in each scanned directory.

    Returns a dict:
      { 0: "<DEVICE_INFO -> DEVICE_NAME>", 1: "<str(directory)> -> <total> = <kept> + <moved> -> <str(dest_dir)>", ... }
    """
    from peepfilt.detector import PersonDetector

    src = Path(input_dir).resolve()
    if not src.exists() or not src.is_dir():
        raise ValueError(f"Input directory not found: {src}")

    # Determine fixed vs per-directory no_people handling
    fixed_no_people_dir = Path(no_people_dir).resolve() if no_people_dir else None
    if fixed_no_people_dir:
        fixed_no_people_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Using fixed no_people directory: {fixed_no_people_dir}")

    detector = PersonDetector()
    device_info = detector.device

    # Device line (key 0)
    device_label = (device_info or "cpu").upper()
    device_name = "CPU"
    try:
        import torch

        if device_info and device_info.lower() == "cuda" and torch.cuda.is_available():
            device_name = torch.cuda.get_device_name(0)
    except Exception:
        pass

    result: Dict[int | str, str] = {0: f"Processed on {device_label} -> {device_name}"}

    # Get directories to scan
    if recursive:
        dir_image_pairs = _iter_images_recursive(src)
        if not dir_image_pairs:
            logger.warning(f"No images found in {src} (recursive)")
            return result
    else:
        images = _iter_images_once(src)
        dir_image_pairs = [(src, images)] if images else []
        if not images:
            logger.warning(f"No images found in {src}")
            return result

    idx = 1

    # Process each directory
    for directory, images in dir_image_pairs:
        logger.info(f"Found {len(images)} images in {directory}")

        # Determine destination for this directory
        if fixed_no_people_dir:
            dest_dir = fixed_no_people_dir
        else:
            dest_dir = directory / "no_people"
            dest_dir.mkdir(parents=True, exist_ok=True)

        per_dir_kept = 0
        per_dir_moved = 0

        for batch in _chunks(images, max(1, int(batch_size))):
            preds = detector.predict_paths(
                [str(p) for p in batch], conf=conf, batch_size=len(batch)
            )
            for p, has_person in zip(batch, preds):
                if has_person:
                    per_dir_kept += 1
                else:
                    shutil.move(str(p), str(dest_dir / p.name))
                    per_dir_moved += 1
                    logger.debug(f"Moved no-person image: {p} -> {dest_dir / p.name}")

        total_processed = per_dir_kept + per_dir_moved
        result[idx] = (
            f"{str(directory)} -> {total_processed} = {per_dir_kept} + {per_dir_moved} -> {str(dest_dir)}"
        )
        idx += 1

    logger.info("Processing complete.")
    return result
