"""
Helpers to select media files for pipeline steps.
"""
import os
from typing import Any, Dict, List, Optional
from fnmatch import fnmatch
from utils.storage import path_exists, read_json


def load_roadmap_file(run_id: str) -> Optional[Dict[str, Any]]:
    """
    Load the roadmap JSON file for the given run ID.
    """
    roadmap_path = f"output/can-do-steps/{run_id}/roadmap-{run_id}.json"

    if not path_exists(roadmap_path):
        print(f"⚠️ Roadmap file not found: {roadmap_path}")
        return None

    try:
        data = read_json(roadmap_path)
        print(f"✓ Loaded roadmap file: {roadmap_path}")
        return data
    except Exception as exc:
        print(f"⚠️ Error loading roadmap file {roadmap_path}: {exc}")
        return None


def _extract_first_image_filename(step: Dict[str, Any]) -> Optional[str]:
    images = step.get("images")
    if isinstance(images, list) and images:
        entry = images[0]
        if isinstance(entry, dict):
            return entry.get("file")
        if isinstance(entry, str):
            return entry
    if isinstance(images, dict):
        return images.get("file")
    if isinstance(images, str):
        return images
    return None


def collect_source_images(source_dir: str) -> List[str]:
    """Collect all candidate image files from a directory."""
    allowed_ext = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
    files: List[str] = []
    for entry in sorted(os.listdir(source_dir)):
        path = os.path.join(source_dir, entry)
        if not os.path.isfile(path):
            continue
        _, ext = os.path.splitext(entry)
        if ext.lower() not in allowed_ext:
            continue
        files.append(path)
    return files


def _filter_files_by_match(files: List[str], match: Optional[str]) -> List[str]:
    if not match:
        return files
    filtered: List[str] = []
    for path in files:
        name = os.path.basename(path)
        if fnmatch(name, match):
            filtered.append(path)
    return filtered


def select_media_files(
    run_id: str,
    source_dir: str,
    match_mode: Optional[str] = None,
    match: Optional[str] = None
) -> List[str]:
    """
    Select media files based on a match mode.

    match_mode values:
      - "first-image": use first image per step from roadmap data
      - None or "all": include all image files in source_dir
    """
    if not match_mode or match_mode == "all":
        return _filter_files_by_match(collect_source_images(source_dir), match)

    if match_mode != "first-image":
        raise ValueError(f"Unsupported match_mode: {match_mode}")

    roadmap_data = load_roadmap_file(run_id)
    if not roadmap_data:
        return []

    filenames: List[str] = []
    seen = set()
    for track in roadmap_data.get("tracks", []):
        for path in track.get("paths", []):
            for step in path.get("steps", []):
                filename = _extract_first_image_filename(step)
                if not filename or filename in seen:
                    continue
                seen.add(filename)
                filenames.append(filename)

    selected = [os.path.join(source_dir, name) for name in filenames]
    return _filter_files_by_match(selected, match)