utils/avatar_manager.py

"""
Avatar Manager Module
=====================
Handles avatar discovery, creation, and management.

Functions:
    - ensure_sample_avatar: Create default sample avatar
    - list_avatars: Get list of available avatars
    - get_avatar_preview: Get preview image of an avatar
"""

from PIL import Image, ImageDraw
from pathlib import Path
from typing import List, Optional
import numpy as np


def ensure_sample_avatar(avatars_dir: Path) -> None:
    """
    Create a sample avatar if none exists.
    
    Generates a simple animated avatar with:
    - Base face image
    - Three mouth positions (closed, medium, open)
    
    Args:
        avatars_dir: Base directory for avatars
        
    Example:
        >>> ensure_sample_avatar(Path("./avatars"))
        # Creates ./avatars/sample/ with base.png and mouth_*.png
    
    Note:
        This creates a basic placeholder avatar. For better results,
        create custom avatars with proper artwork.
    """
    sample_dir = avatars_dir / "sample"
    
    # Check if sample already exists with content
    if sample_dir.exists() and any(sample_dir.iterdir()):
        return
    
    # Create directory
    sample_dir.mkdir(parents=True, exist_ok=True)
    
    # Image dimensions
    width, height = 512, 512
    
    # Create base image (simple face background)
    base = Image.new("RGBA", (width, height), (255, 220, 200, 255))
    draw_base = ImageDraw.Draw(base)
    
    # Draw simple face features on base
    # Face circle
    draw_base.ellipse([56, 56, 456, 456], fill=(255, 230, 210, 255), outline=(200, 150, 130, 255), width=3)
    
    # Eyes
    draw_base.ellipse([150, 180, 200, 230], fill=(255, 255, 255, 255), outline=(0, 0, 0, 255), width=2)
    draw_base.ellipse([312, 180, 362, 230], fill=(255, 255, 255, 255), outline=(0, 0, 0, 255), width=2)
    
    # Pupils
    draw_base.ellipse([165, 195, 185, 215], fill=(50, 50, 50, 255))
    draw_base.ellipse([327, 195, 347, 215], fill=(50, 50, 50, 255))
    
    # Eyebrows
    draw_base.arc([140, 150, 210, 190], start=200, end=340, fill=(100, 70, 50, 255), width=3)
    draw_base.arc([302, 150, 372, 190], start=200, end=340, fill=(100, 70, 50, 255), width=3)
    
    # Nose
    draw_base.polygon([(256, 250), (240, 310), (272, 310)], fill=(240, 200, 180, 255))
    
    # Hair (simple)
    draw_base.arc([40, 20, 472, 300], start=180, end=360, fill=(80, 50, 30, 255), width=30)
    
    base.save(sample_dir / "base.png")
    
    # Create mouth frames (transparent overlays)
    mouth_positions = [
        # (y_offset, height) - Mouth closed to open
        (0, 8),    # mouth_0: Nearly closed
        (0, 20),   # mouth_1: Slightly open
        (0, 35),   # mouth_2: Wide open
    ]
    
    for i, (y_off, mouth_height) in enumerate(mouth_positions):
        # Create transparent image for mouth overlay
        mouth_img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
        draw_mouth = ImageDraw.Draw(mouth_img)
        
        # Calculate mouth position
        mouth_y = 340 + y_off
        mouth_left = 200
        mouth_right = 312
        
        # Draw mouth (ellipse shape)
        draw_mouth.ellipse(
            [mouth_left, mouth_y, mouth_right, mouth_y + mouth_height],
            fill=(180, 80, 80, 255),
            outline=(120, 50, 50, 255),
            width=2
        )
        
        # Add inner mouth detail for open mouths
        if mouth_height > 15:
            inner_offset = 5
            draw_mouth.ellipse(
                [mouth_left + inner_offset, mouth_y + inner_offset,
                 mouth_right - inner_offset, mouth_y + mouth_height - inner_offset],
                fill=(100, 40, 40, 255)
            )
        
        mouth_img.save(sample_dir / f"mouth_{i}.png")


def list_avatars(avatars_dir: Path) -> List[str]:
    """
    Get list of available avatar names.
    
    Scans the avatars directory for valid avatar folders
    (containing base.png and mouth_*.png files).
    
    Args:
        avatars_dir: Base directory containing avatar folders
        
    Returns:
        List of avatar folder names
        
    Example:
        >>> avatars = list_avatars(Path("./avatars"))
        >>> print(avatars)
        ['sample', 'anime_girl', 'anime_boy']
    """
    # Ensure sample avatar exists
    ensure_sample_avatar(avatars_dir)
    
    # Find all valid avatar directories
    avatars = []
    
    if avatars_dir.exists():
        for path in avatars_dir.iterdir():
            if path.is_dir():
                # Check for required files
                has_base = (path / "base.png").exists()
                has_mouth = any(path.glob("mouth_*.png"))
                
                if has_base and has_mouth:
                    avatars.append(path.name)
    
    return sorted(avatars)


def get_avatar_preview(avatar_name: str, avatars_dir: Path) -> Optional[Image.Image]:
    """
    Get a preview image of an avatar.
    
    Composites the base image with the first mouth frame
    to show what the avatar looks like.
    
    Args:
        avatar_name: Name of the avatar folder
        avatars_dir: Base directory containing avatar folders
        
    Returns:
        PIL Image object or None if avatar not found
        
    Example:
        >>> preview = get_avatar_preview("sample", Path("./avatars"))
        >>> preview.show()
    """
    avatar_folder = avatars_dir / avatar_name
    base_path = avatar_folder / "base.png"
    
    if not base_path.exists():
        return None
    
    # Load base image
    base = Image.open(base_path).convert("RGBA")
    
    # Find first mouth frame
    mouth_frames = sorted(avatar_folder.glob("mouth_*.png"))
    
    if mouth_frames:
        mouth = Image.open(mouth_frames[0]).convert("RGBA").resize(base.size)
        # Composite mouth onto base
        preview = Image.alpha_composite(base, mouth)
    else:
        preview = base
    
    return preview