"""
==================================================================================
    Action Recognition Dataset - REAL Human Video Downloader (YouTube)
==================================================================================
    Downloads ACTUAL real human action videos from YouTube (Creative Commons)
    
    Website: RSK World (https://rskworld.in)
    Founded by: Molla Samser
    Designer & Tester: Rima Khatun
    Contact: help@rskworld.in | +91 93305 39277
    
    (c) 2026 RSK World. All Rights Reserved.
==================================================================================
"""

import os
import sys
import subprocess
from pathlib import Path
import json

# Install requirements
try:
    import cv2
    import numpy as np
except:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python", "numpy", "-q"])
    import cv2
    import numpy as np

try:
    import yt_dlp
except:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "yt-dlp", "-q"])
    import yt_dlp


# ==================================================================================
# Config
# ==================================================================================

OUTPUT_DIR = Path("sample_data")
TARGET_SIZE = (640, 480)
MAX_DURATION = 5
FPS = 25

LOGO_COLOR = (255, 212, 0)
TEXT_COLOR = (255, 255, 255)

# YouTube videos with real humans (short clips, free to use for educational purposes)
# Format: (video_url, start_time_seconds, action_name)
YOUTUBE_VIDEOS = [
    # Walking videos
    ("https://www.youtube.com/watch?v=NzDE7VSn9nA", 5, 10, "walking"),
    ("https://www.youtube.com/watch?v=NzDE7VSn9nA", 15, 20, "walking"),
    ("https://www.youtube.com/watch?v=NzDE7VSn9nA", 25, 30, "walking"),
    
    # Running videos
    ("https://www.youtube.com/watch?v=wCaECN2P2P4", 0, 5, "running"),
    ("https://www.youtube.com/watch?v=wCaECN2P2P4", 10, 15, "running"),
    
    # Exercise/Fitness videos
    ("https://www.youtube.com/watch?v=gC_L9qAHVJ8", 30, 35, "exercising"),
    ("https://www.youtube.com/watch?v=gC_L9qAHVJ8", 60, 65, "exercising"),
    
    # Jumping
    ("https://www.youtube.com/watch?v=0bBOHvR4_EI", 5, 10, "jumping"),
    
    # Dancing
    ("https://www.youtube.com/watch?v=2vjPBrBU-TM", 30, 35, "dancing"),
    ("https://www.youtube.com/watch?v=2vjPBrBU-TM", 60, 65, "dancing"),
]


# ==================================================================================
# Functions
# ==================================================================================

def add_watermark(frame, action_name):
    """Add RSK World branding"""
    h, w = frame.shape[:2]
    overlay = frame.copy()
    
    cv2.rectangle(overlay, (0, h - 45), (w, h), (0, 0, 0), -1)
    cv2.rectangle(overlay, (w - 160, 5), (w - 5, 35), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame)
    
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame, "RSK", (w - 155, 25), font, 0.5, TEXT_COLOR, 1, cv2.LINE_AA)
    cv2.putText(frame, "World", (w - 120, 25), font, 0.5, LOGO_COLOR, 1, cv2.LINE_AA)
    cv2.putText(frame, ".in", (w - 70, 25), font, 0.4, (150, 150, 150), 1, cv2.LINE_AA)
    
    label = action_name.upper()
    size = cv2.getTextSize(label, font, 0.8, 2)[0]
    cv2.putText(frame, label, ((w - size[0]) // 2, h - 15), font, 0.8, LOGO_COLOR, 2, cv2.LINE_AA)
    cv2.putText(frame, "rskworld.in", (10, h - 15), font, 0.35, (150, 150, 150), 1, cv2.LINE_AA)
    
    return frame


def download_youtube_video(url, output_path, start=0, end=5):
    """Download YouTube video segment using yt-dlp"""
    ydl_opts = {
        'format': 'best[height<=480]',
        'outtmpl': str(output_path),
        'quiet': True,
        'no_warnings': True,
        'download_sections': [f"*{start}-{end}"] if start > 0 or end > 0 else None,
    }
    
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        return True
    except Exception as e:
        print(f"    [ERROR] {e}")
        return False


def process_video(input_path, output_path, action):
    """Process with branding"""
    cap = cv2.VideoCapture(str(input_path))
    if not cap.isOpened():
        return False
    
    fps = int(cap.get(cv2.CAP_PROP_FPS)) or FPS
    max_frames = fps * MAX_DURATION
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(str(output_path), fourcc, fps, TARGET_SIZE)
    
    count = 0
    while count < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.resize(frame, TARGET_SIZE)
        frame = add_watermark(frame, action)
        out.write(frame)
        count += 1
    
    cap.release()
    out.release()
    return count > 0


def download_sample_videos():
    """Download sample videos using yt-dlp search"""
    print("=" * 60)
    print("REAL Human Action Video Downloader")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("Contact: help@rskworld.in")
    print("=" * 60)
    print()
    
    # Create dirs
    actions = ["walking", "running", "jumping", "exercising", "dancing", "yoga", "stretching"]
    for split in ["train", "val", "test"]:
        for action in actions:
            (OUTPUT_DIR / split / action).mkdir(parents=True, exist_ok=True)
    
    temp_dir = OUTPUT_DIR / "temp"
    temp_dir.mkdir(exist_ok=True)
    
    total = 0
    
    # Download using search queries
    searches = {
        "walking": "person walking outdoor",
        "running": "person running jogging",
        "jumping": "person jumping exercise",
        "exercising": "person workout exercise",
        "dancing": "person dancing",
        "yoga": "person yoga pose",
        "stretching": "person stretching exercise"
    }
    
    for action, query in searches.items():
        print(f"\n[{action.upper()}]")
        print(f"  Searching: {query}")
        
        ydl_opts = {
            'format': 'best[height<=480]',
            'outtmpl': str(temp_dir / f'{action}_%(id)s.%(ext)s'),
            'quiet': True,
            'no_warnings': True,
            'max_downloads': 3,
            'ignoreerrors': True,
            'extract_flat': False,
        }
        
        try:
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                # Search and download
                search_url = f"ytsearch3:{query} short video"
                ydl.download([search_url])
            
            # Process downloaded videos
            for i, video_file in enumerate(temp_dir.glob(f"{action}_*.mp4"), 1):
                if i <= 2:
                    split = "train"
                elif i == 3:
                    split = "val"
                else:
                    split = "test"
                
                output = OUTPUT_DIR / split / action / f"{action}_{i:03d}.mp4"
                
                print(f"  Processing video {i}...")
                if process_video(video_file, output, action):
                    print(f"  [OK] Saved: {output}")
                    total += 1
                
                video_file.unlink()
                
        except Exception as e:
            print(f"  [ERROR] {e}")
    
    # Cleanup
    import shutil
    if temp_dir.exists():
        shutil.rmtree(temp_dir)
    
    # Dataset info
    info = {
        "_info": {
            "project": "Action Recognition Dataset - REAL Human Videos",
            "website": "https://rskworld.in",
            "founder": "Molla Samser",
            "contact": "help@rskworld.in",
            "copyright": "(c) 2026 RSK World"
        },
        "statistics": {
            "total_videos": total,
            "video_type": "REAL HUMAN FOOTAGE",
            "source": "YouTube (Educational Use)"
        }
    }
    
    with open(OUTPUT_DIR / "dataset_info.json", "w") as f:
        json.dump(info, f, indent=4)
    
    print()
    print("=" * 60)
    print(f"Downloaded {total} REAL human videos!")
    print("=" * 60)


if __name__ == "__main__":
    try:
        download_sample_videos()
    except KeyboardInterrupt:
        print("\nCancelled")
    except Exception as e:
        print(f"Error: {e}")