"""
Process downloaded YouTube videos with RSK World branding
RSK World (https://rskworld.in) | Founder: Molla Samser
(c) 2026 RSK World. All Rights Reserved.
"""
import cv2
import numpy as np
from pathlib import Path
import json

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

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

def add_watermark(frame, action_name):
    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):
    print(f"  Processing: {input_path.name} -> {output_path}")
    
    cap = cv2.VideoCapture(str(input_path))
    if not cap.isOpened():
        print(f"    [ERROR] Cannot open {input_path}")
        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()
    
    print(f"    [OK] Created with {count} frames")
    return count > 0

def main():
    print("=" * 60)
    print("Processing REAL Human Videos with RSK World Branding")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("=" * 60)
    print()
    
    # Find all temp files
    temp_files = list(OUTPUT_DIR.glob("temp_*"))
    print(f"Found {len(temp_files)} downloaded videos")
    print()
    
    processed = 0
    
    for temp_file in temp_files:
        # Extract action name and number from filename
        # Format: temp_action_001
        parts = temp_file.stem.split("_")
        if len(parts) >= 3:
            action = parts[1]
            num = parts[2]
        else:
            continue
        
        # Determine split based on number
        num_int = int(num)
        if num_int == 1:
            split = "train"
        elif num_int == 2:
            split = "val" 
        else:
            split = "test"
        
        # Create output directory
        out_dir = OUTPUT_DIR / split / action
        out_dir.mkdir(parents=True, exist_ok=True)
        
        output_path = out_dir / f"{action}_{num}.mp4"
        
        print(f"\n[{action.upper()} - {split}]")
        if process_video(temp_file, output_path, action):
            processed += 1
            # Delete temp file after successful processing
            temp_file.unlink()
    
    # Create dataset info
    info = {
        "_info": {
            "project": "Action Recognition Dataset",
            "website": "https://rskworld.in",
            "founder": "Molla Samser",
            "designer": "Rima Khatun",
            "contact": "help@rskworld.in",
            "phone": "+91 93305 39277",
            "copyright": "(c) 2026 RSK World"
        },
        "source": "YouTube Real Human Videos",
        "statistics": {
            "processed": processed,
            "resolution": "640x480",
            "fps": 30,
            "max_duration": "5 seconds"
        }
    }
    
    with open(OUTPUT_DIR / "dataset_info.json", "w") as f:
        json.dump(info, f, indent=4)
    
    print()
    print("=" * 60)
    print(f"Complete! Processed {processed} REAL human videos!")
    print("=" * 60)
    print()
    print("All videos include RSK World branding!")
    print("Visit: https://rskworld.in")

if __name__ == "__main__":
    main()