"""
Image Management Module

Responsible for indexing, validating, and selecting facial images required for the game.
"""

import os
import random
from pathlib import Path
from typing import List, Dict, Literal


class ImagePool:
    """Manages indexing and selection of AI-generated and real human face images"""

    def __init__(self, ai_dir: str = "image/ai", real_dir: str = "image/real"):
        """
        Initialize image pool

        Args:
            ai_dir: Directory path for AI-generated images
            real_dir: Directory path for real human images

        Raises:
            ValueError: If either directory contains less than 10 images
        """
        self.ai_dir = ai_dir
        self.real_dir = real_dir
        self.ai_images = self._scan_directory(ai_dir)
        self.real_images = self._scan_directory(real_dir)
        self.validate()

    def _scan_directory(self, directory: str) -> List[str]:
        """
        Scan directory and return paths of all valid image files

        Args:
            directory: Directory path to scan

        Returns:
            List of image file paths
        """
        valid_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
        image_paths = []

        if not os.path.exists(directory):
            return image_paths

        for file in Path(directory).rglob('*'):
            if file.is_file() and file.suffix.lower() in valid_extensions:
                image_paths.append(str(file))

        return sorted(image_paths)  # Sort for consistency

    @property
    def total_ai(self) -> int:
        """Return total number of AI-generated images"""
        return len(self.ai_images)

    @property
    def total_real(self) -> int:
        """Return total number of real human images"""
        return len(self.real_images)

    def validate(self) -> None:
        """
        Validate that image count meets requirements

        Raises:
            ValueError: If either category has less than 10 images
        """
        if self.total_ai < 10:
            raise ValueError(
                f"Insufficient AI images: need at least 10, found {self.total_ai} in {self.ai_dir}"
            )
        if self.total_real < 10:
            raise ValueError(
                f"Insufficient real images: need at least 10, found {self.total_real} in {self.real_dir}"
            )

    def create_game_set(self) -> List[Dict[str, str]]:
        """
        Create a random set of 20 images for a new game (10 AI + 10 real)

        Returns:
            List of 20 ImageItem dictionaries, each containing 'path' and 'label' fields
            Images are in random order
        """
        # Randomly select 10 AI images
        selected_ai = random.sample(self.ai_images, 10)

        # Randomly select 10 real images
        selected_real = random.sample(self.real_images, 10)

        # Create ImageItem dictionary list
        all_images = [
            {"path": path, "label": "AI"} for path in selected_ai
        ] + [
            {"path": path, "label": "Human"} for path in selected_real
        ]

        # Shuffle order
        random.shuffle(all_images)

        return all_images

    def get_random_image(self, label: Literal["AI", "Human"]) -> str:
        """
        Get a random image path from specified category

        Args:
            label: Image category, "AI" or "Human"

        Returns:
            Randomly selected image path
        """
        if label == "AI":
            return random.choice(self.ai_images)
        else:
            return random.choice(self.real_images)