"""
Scene Image Selector - Randomly selects candidate images for each scene
Saves selections to a folder before composition
"""
import os
import json
import random
import shutil
from pathlib import Path


class SceneSelector:
    def __init__(self, manifest_path, candidates_dir, output_dir="selected"):
        """
        Initialize the scene selector.
        
        Args:
            manifest_path: Path to manifest_response.json
            candidates_dir: Path to candidates folder with scene_0-6 subfolders
            output_dir: Where to save selected images
        """
        self.manifest_path = manifest_path
        self.candidates_dir = candidates_dir
        self.output_dir = output_dir
        self.selections = {}
        
        # Create output directory
        Path(self.output_dir).mkdir(parents=True, exist_ok=True)
        
        # Load manifest
        with open(manifest_path, 'r') as f:
            self.manifest = json.load(f)
        
        print(f"[Selector] Initialized with {len(self.manifest.get('scenes', []))} scenes")
        print(f"[Selector] Candidates dir: {candidates_dir}")
        print(f"[Selector] Output dir: {output_dir}\n")

    def select_all_scenes(self):
        """Randomly select one image per scene."""
        scenes = self.manifest.get("scenes", [])
        print(f"[Selector] Selecting images for {len(scenes)} scenes...\n")
        
        for scene_idx, scene in enumerate(scenes):
            self._select_scene(scene_idx, scene)
        
        # Save selections metadata
        self._save_selections_metadata()
        
        print(f"\n[Selector] ✓ Selected {len(scenes)} images!")
        print(f"[Selector] Check {self.output_dir}/ for selected images")
        print(f"[Selector] Check {self.output_dir}/selections.json for metadata")

    def _select_scene(self, scene_idx, scene):
        """Randomly select one image for a scene."""
        scene_label = scene.get("label", f"Scene {scene_idx}")
        candidates_folder = os.path.join(self.candidates_dir, f"scene_{scene_idx}")
        
        # Get list of candidate images
        if not os.path.exists(candidates_folder):
            print(f"[Scene {scene_idx}] ✗ Folder not found: {candidates_folder}")
            return
        
        candidates = [f for f in os.listdir(candidates_folder) if f.endswith(('.jpg', '.png'))]
        if not candidates:
            print(f"[Scene {scene_idx}] ✗ No images found in {candidates_folder}")
            return
        
        # Randomly select one image
        selected_image = random.choice(candidates)
        source_path = os.path.join(candidates_folder, selected_image)
        
        # Copy to output folder
        output_filename = f"scene_{scene_idx:02d}.jpg"
        dest_path = os.path.join(self.output_dir, output_filename)
        shutil.copy2(source_path, dest_path)
        
        # Store selection metadata
        self.selections[scene_idx] = {
            "scene_label": scene_label,
            "original_file": selected_image,
            "source_folder": candidates_folder,
            "total_candidates": len(candidates),
            "output_file": output_filename
        }
        
        print(f"[Scene {scene_idx}] {scene_label}")
        print(f"  Available: {len(candidates)} candidates")
        print(f"  Selected: {selected_image}")
        print(f"  Copied to: {output_filename}\n")

    def _save_selections_metadata(self):
        """Save metadata about which images were selected."""
        metadata_path = os.path.join(self.output_dir, "selections.json")
        with open(metadata_path, 'w') as f:
            json.dump(self.selections, f, indent=2)
        print(f"[Selector] Saved selections metadata to {metadata_path}")


if __name__ == "__main__":
    # Load manifest from content_gen_server
    manifest_file = "../content_gen_server/manifest_response.json"
    candidates_dir = "../workspace/renders/candidates"
    output_dir = "selected"
    
    # Create and run selector
    selector = SceneSelector(manifest_file, candidates_dir, output_dir)
    selector.select_all_scenes()