modules/text_story/services/background.py

"""
Background Handler for Text Story module.
Handles gameplay video loading from HuggingFace Dataset storage.
"""

import os
import random
import logging
from moviepy.editor import VideoFileClip, vfx
from typing import Optional, List
from huggingface_hub import hf_hub_download, list_repo_files

logger = logging.getLogger(__name__)

# Folder name in HF Dataset for gameplay backgrounds
HF_BACKGROUNDS_FOLDER = "gameplay_backgrounds"

# Local cache path
LOCAL_CACHE_DIR = "cache/gameplay_backgrounds"


class BackgroundHandler:
    """
    Handles gameplay background video processing.
    Downloads from HuggingFace Dataset (HF_REPO env variable).
    
    Features:
    - Download videos from HF Dataset
    - Random video selection
    - Audio removal
    - Slow motion (0.7x)
    - Dark overlay
    - Seamless looping
    """
    
    def __init__(self):
        # Get repo from environment variable (e.g., robiul487/NCAkit)
        self.repo_id = os.getenv("HF_REPO", "")
        self.folder = HF_BACKGROUNDS_FOLDER
        self.cache_dir = LOCAL_CACHE_DIR
        
        if not self.repo_id:
            logger.warning("BackgroundHandler: HF_REPO not set! Using green screen background.")
            self.available_videos = []
            return
        
        # Ensure cache directory exists
        os.makedirs(self.cache_dir, exist_ok=True)
        
        # Get list of available videos
        self.available_videos = self._list_available_videos()
        
        if self.available_videos:
            logger.info(f"BackgroundHandler: Found {len(self.available_videos)} videos in {self.repo_id}/{self.folder}")
        else:
            logger.warning(f"BackgroundHandler: No videos found in {self.repo_id}/{self.folder}")
            # Try to create the folder in HF Dataset
            self._ensure_folder_exists()
    
    def _list_available_videos(self) -> List[str]:
        """List available video files in HF Dataset folder."""
        if not self.repo_id:
            return []
        
        try:
            all_files = list_repo_files(
                repo_id=self.repo_id,
                repo_type="dataset"
            )
            
            # Filter for .mp4 videos in gameplay_backgrounds folder only
            videos = [
                f for f in all_files 
                if f.startswith(f"{self.folder}/") 
                and f.lower().endswith('.mp4')
            ]
            
            return videos
            
        except Exception as e:
            logger.error(f"BackgroundHandler: Failed to list files - {e}")
            return []
    
    def _ensure_folder_exists(self):
        """Create gameplay_backgrounds folder in HF Dataset with a placeholder file."""
        try:
            from huggingface_hub import HfApi
            import tempfile
            
            hf_token = os.getenv("HF_TOKEN", "")
            if not hf_token:
                logger.warning("BackgroundHandler: HF_TOKEN not set, cannot create folder")
                return
            
            api = HfApi(token=hf_token)
            
            # Create a README placeholder file to create the folder
            placeholder_content = """# Gameplay Backgrounds

Place your background videos here (.mp4 only).

These videos will be:
- Randomly selected for text story backgrounds
- Slowed down (0.7x)
- Darkened for better text visibility
- Center-cropped if 16:9

Recommended: 9:16 vertical videos (TikTok/Reels format)
"""
            
            # Upload placeholder to create the folder
            with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
                f.write(placeholder_content)
                temp_path = f.name
            
            api.upload_file(
                path_or_fileobj=temp_path,
                path_in_repo=f"{self.folder}/README.md",
                repo_id=self.repo_id,
                repo_type="dataset"
            )
            
            # Cleanup temp file
            os.remove(temp_path)
            
            logger.info(f"BackgroundHandler: Created folder {self.folder}/ in HF Dataset")
            logger.info(f"BackgroundHandler: Upload .mp4 videos to {self.repo_id}/{self.folder}/")
            
        except Exception as e:
            logger.error(f"BackgroundHandler: Failed to create folder - {e}")
    
    def _download_video(self, filename: str) -> Optional[str]:
        """Download a video from HF Dataset to local cache."""
        try:
            # Check if already cached
            local_name = os.path.basename(filename)
            cached_path = os.path.join(self.cache_dir, local_name)
            
            if os.path.exists(cached_path):
                logger.info(f"BackgroundHandler: Using cached {local_name}")
                return cached_path
            
            # Download from HF
            logger.info(f"BackgroundHandler: Downloading {filename}...")
            
            downloaded_path = hf_hub_download(
                repo_id=self.repo_id,
                filename=filename,
                repo_type="dataset",
                local_dir=self.cache_dir,
                local_dir_use_symlinks=False
            )
            
            logger.info(f"BackgroundHandler: Downloaded to {downloaded_path}")
            return downloaded_path
            
        except Exception as e:
            logger.error(f"BackgroundHandler: Download failed - {e}")
            return None
    
    def get_random_video(self) -> Optional[str]:
        """Get random video from HF Dataset and download it."""
        if not self.available_videos:
            # Retry listing
            self.available_videos = self._list_available_videos()
        
        if not self.available_videos:
            logger.warning("BackgroundHandler: No videos available")
            return None
        
        selected = random.choice(self.available_videos)
        logger.info(f"BackgroundHandler: Selected {selected}")
        
        return self._download_video(selected)
    
    def load_and_process(self, 
                         target_duration: float,
                         video_path: str = None) -> Optional[VideoFileClip]:
        """
        Load and process a background video.
        
        Args:
            target_duration: Required duration in seconds
            video_path: Optional specific video path (or random if None)
            
        Returns:
            Processed VideoFileClip or None
        """
        # Get video path
        if video_path is None:
            video_path = self.get_random_video()
        
        if not video_path or not os.path.exists(video_path):
            logger.warning("BackgroundHandler: No video available, creating solid background")
            return self._create_solid_background(target_duration)
        
        try:
            # Load video WITHOUT audio
            clip = VideoFileClip(video_path).without_audio()
            logger.info(f"BackgroundHandler: Loaded {video_path}, duration: {clip.duration:.1f}s")
            
            # Apply slow motion (0.7x speed)
            clip = clip.fx(vfx.speedx, 0.7)
            
            # Loop if needed to match target duration
            clip = self._loop_to_duration(clip, target_duration)
            
            # Apply visual effects
            clip = self._apply_visual_effects(clip)
            
            return clip
            
        except Exception as e:
            logger.error(f"BackgroundHandler: Failed to process video - {e}")
            return self._create_solid_background(target_duration)
    
    def _loop_to_duration(self, clip: VideoFileClip, target_duration: float) -> VideoFileClip:
        """Loop video to match target duration."""
        if clip.duration >= target_duration:
            return clip.subclip(0, target_duration)
        
        loops_needed = int(target_duration / clip.duration) + 1
        looped = clip.loop(n=loops_needed)
        return looped.subclip(0, target_duration)
    
    def _apply_visual_effects(self, clip: VideoFileClip) -> VideoFileClip:
        """
        Apply resize, crop (if needed), dark overlay, and saturation reduction.
        
        - 9:16 videos: just resize (no crop needed)
        - 16:9 videos: center crop to 9:16
        """
        target_w, target_h = 1080, 1920
        target_ratio = target_w / target_h  # 0.5625 (9:16)
        
        # Calculate source aspect ratio
        clip_ratio = clip.w / clip.h
        
        # Check if already 9:16 (or close to it)
        # 9:16 ratio is ~0.5625, allow some tolerance
        is_vertical = clip_ratio < 0.7  # Less than ~11:16 is considered vertical
        
        if is_vertical:
            # Already vertical (9:16)
            if clip.w == target_w and clip.h == target_h:
                # Perfect match, no resize needed
                logger.info(f"BackgroundHandler: Video is already {target_w}x{target_h}, no resize")
            else:
                # Resize to target resolution
                logger.info(f"BackgroundHandler: Video is vertical ({clip.w}x{clip.h}), resizing to {target_w}x{target_h}")
                clip = clip.resize(newsize=(target_w, target_h))
        else:
            # Horizontal (16:9), need to crop
            logger.info(f"BackgroundHandler: Video is horizontal ({clip.w}x{clip.h}), center cropping to 9:16")
            
            # Scale to match height, then center crop width
            new_h = target_h
            new_w = int(clip_ratio * new_h)
            clip = clip.resize(height=new_h)
            
            # Center crop
            x_center = new_w // 2
            clip = clip.crop(x_center=x_center, width=target_w, height=target_h)
        
        # Dark overlay (reduce brightness by 40%)
        clip = clip.fx(vfx.colorx, 0.6)
        
        # Saturation reduction
        clip = clip.fx(vfx.lum_contrast, lum=-10, contrast=-0.1)
        
        return clip
    
    def _create_solid_background(self, duration: float) -> VideoFileClip:
        """Create GREEN SCREEN background for chroma key overlay."""
        from moviepy.editor import ColorClip
        
        # Green screen color (0, 255, 0) for easy chroma key removal
        return ColorClip(
            size=(1080, 1920),
            color=(0, 255, 0),  # Pure green for chroma key
            duration=duration
        )