"""
==================================================================================
    Action Recognition Dataset - YouTube Real Human Video Downloader
==================================================================================
    Project: Action Recognition Dataset
    
    Downloads REAL human action videos from YouTube using yt-dlp
    
==================================================================================
    DEVELOPER INFORMATION
==================================================================================
    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 dependencies
def install_deps():
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "yt-dlp", "opencv-python", "numpy"])

try:
    import cv2
    import numpy as np
    import yt_dlp
except ImportError:
    print("Installing dependencies...")
    install_deps()
    import cv2
    import numpy as np
    import yt_dlp

# ==================================================================================
# Configuration
# ==================================================================================

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

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

# YouTube video IDs for real human action videos (short clips)
# Using search queries to find appropriate videos
YOUTUBE_SEARCHES = {
    "walking": "person walking street short clip",
    "running": "person running jogging short clip",
    "jumping": "person jumping short clip",
    "dancing": "person dancing short clip",
    "exercising": "person workout exercise short clip",
    "yoga": "yoga pose short clip",
    "stretching": "stretching exercise short clip",
    "boxing": "boxing training short clip"
}

# Specific video IDs (public/creative commons when possible)
YOUTUBE_VIDEOS = {
    "walking": [
        {"url": "https://www.youtube.com/shorts/Q6vXGKCv-Uk", "name": "walking_001"},
        {"url": "https://www.youtube.com/shorts/nJvz3g7GHDY", "name": "walking_002"},
    ],
    "running": [
        {"url": "https://www.youtube.com/shorts/jM7dCa1IYPM", "name": "running_001"},
        {"url": "https://www.youtube.com/shorts/aCNDdX1gXoU", "name": "running_002"},
    ],
    "jumping": [
        {"url": "https://www.youtube.com/shorts/z6nLHPj42Kg", "name": "jumping_001"},
    ],
    "dancing": [
        {"url": "https://www.youtube.com/shorts/EcNf57EXbpk", "name": "dancing_001"},
        {"url": "https://www.youtube.com/shorts/iB1uQBzE1gc", "name": "dancing_002"},
    ],
    "exercising": [
        {"url": "https://www.youtube.com/shorts/BSb7F4aXB5I", "name": "exercising_001"},
        {"url": "https://www.youtube.com/shorts/1f8yoFFdkcY", "name": "exercising_002"},
    ]
}


# ==================================================================================
# RSK World Branding
# ==================================================================================

def add_watermark(frame, action_name):
    """Add RSK World watermark"""
    h, w = frame.shape[:2]
    overlay = frame.copy()
    
    cv2.rectangle(overlay, (0, h - 50), (w, h), (0, 0, 0), -1)
    cv2.rectangle(overlay, (w - 190, 5), (w - 5, 40), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame)
    
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame, "RSK", (w - 180, 30), font, 0.7, TEXT_COLOR, 2, cv2.LINE_AA)
    cv2.putText(frame, "World", (w - 130, 30), font, 0.7, LOGO_COLOR, 2, cv2.LINE_AA)
    cv2.putText(frame, ".in", (w - 60, 30), font, 0.5, (150, 150, 150), 1, cv2.LINE_AA)
    
    label = action_name.upper()
    size = cv2.getTextSize(label, font, 1.0, 2)[0]
    cv2.putText(frame, label, ((w - size[0]) // 2, h - 15), font, 1.0, LOGO_COLOR, 2, cv2.LINE_AA)
    cv2.putText(frame, "rskworld.in", (10, h - 15), font, 0.4, (150, 150, 150), 1, cv2.LINE_AA)
    
    return frame


def process_video(input_path, output_path, action_name):
    """Process downloaded video with branding"""
    cap = cv2.VideoCapture(str(input_path))
    if not cap.isOpened():
        return False
    
    fps = int(cap.get(cv2.CAP_PROP_FPS)) or FPS
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    max_frames = min(total, fps * MAX_DURATION)
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(str(output_path), fourcc, min(fps, 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_name)
        out.write(frame)
        count += 1
    
    cap.release()
    out.release()
    return count > 0


# ==================================================================================
# YouTube Download
# ==================================================================================

def download_youtube_video(url, output_path):
    """Download video from YouTube using yt-dlp"""
    ydl_opts = {
        'format': 'best[height<=480]/best',
        'outtmpl': str(output_path),
        'quiet': True,
        'no_warnings': True,
        'extract_flat': False,
        'socket_timeout': 30,
    }
    
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        return True
    except Exception as e:
        print(f"      [ERROR] {str(e)[:50]}")
        return False


def search_and_download(query, output_path, max_results=1):
    """Search YouTube and download first result"""
    ydl_opts = {
        'format': 'best[height<=480]/best',
        'outtmpl': str(output_path),
        'quiet': True,
        'no_warnings': True,
        'default_search': 'ytsearch',
        'max_downloads': max_results,
        'socket_timeout': 30,
        'match_filter': lambda info: None if info.get('duration', 0) < 120 else 'Video too long'
    }
    
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([f"ytsearch{max_results}:{query}"])
        return True
    except Exception as e:
        print(f"      [ERROR] {str(e)[:50]}")
        return False


# ==================================================================================
# Main
# ==================================================================================

def create_directories():
    """Create output directories"""
    actions = list(YOUTUBE_SEARCHES.keys())
    for split in ["train", "val", "test"]:
        for action in actions:
            (OUTPUT_DIR / split / action).mkdir(parents=True, exist_ok=True)
    print("[OK] Directories created")


def download_all():
    """
    Download real human videos from YouTube
    
    RSK World (https://rskworld.in)
    """
    print("=" * 60)
    print("Action Recognition - REAL Human Video Downloader")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("Designer: Rima Khatun")
    print("Contact: help@rskworld.in | +91 93305 39277")
    print("=" * 60)
    print()
    print("Downloading REAL human action videos from YouTube...")
    print()
    
    create_directories()
    print()
    
    total = 0
    failed = 0
    
    # Download from search queries
    for action, query in YOUTUBE_SEARCHES.items():
        print(f"\n[{action.upper()}]")
        
        for i in range(3):  # 3 videos per action
            split = "train" if i < 1 else ("val" if i == 1 else "test")
            name = f"{action}_{i+1:03d}"
            
            temp_path = OUTPUT_DIR / f"temp_{name}"
            output_path = OUTPUT_DIR / split / action / f"{name}.mp4"
            
            print(f"  [{split}] {name}")
            print(f"      Searching: '{query}'...")
            
            # Search for video
            search_query = f"{query} short"
            if search_and_download(search_query, temp_path, max_results=1):
                # Find the downloaded file
                temp_files = list(OUTPUT_DIR.glob(f"temp_{name}*"))
                if temp_files:
                    temp_file = temp_files[0]
                    if process_video(temp_file, output_path, action):
                        print(f"      [OK] Saved: {output_path}")
                        total += 1
                    else:
                        print(f"      [ERROR] Processing failed")
                        failed += 1
                    # Cleanup
                    temp_file.unlink()
                else:
                    print(f"      [ERROR] Download file not found")
                    failed += 1
            else:
                failed += 1
    
    # Save dataset info
    info = {
        "_info": {
            "project": "Action Recognition Dataset",
            "website": "https://rskworld.in",
            "founder": "Molla Samser",
            "designer": "Rima Khatun",
            "contact": "help@rskworld.in",
            "copyright": "(c) 2026 RSK World"
        },
        "source": "YouTube (via yt-dlp)",
        "note": "Real human action videos",
        "statistics": {
            "downloaded": total,
            "failed": failed
        }
    }
    
    with open(OUTPUT_DIR / "dataset_info.json", "w") as f:
        json.dump(info, f, indent=4)
    
    print()
    print("=" * 60)
    print(f"Complete! Downloaded: {total}, Failed: {failed}")
    print("=" * 60)
    print()
    print("All videos have RSK World branding!")
    print("Visit: https://rskworld.in")


if __name__ == "__main__":
    try:
        download_all()
    except KeyboardInterrupt:
        print("\n\nCancelled.")
    except Exception as e:
        print(f"\nError: {e}")
        print("Contact: help@rskworld.in")