import os
import subprocess
from yt_dlp import YoutubeDL
from .utils import clip_video, get_url_list, get_video_duration, timestamp_to_seconds

def download_videos(url_list, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    for idx, item in enumerate(url_list, 1):
        url = item['url']
        task = item['task']
        index_str = f"{idx:03d}"
        
        existing_files = [f for f in os.listdir(output_dir) if f.startswith(f"{index_str}_")]
        if existing_files:
            print(f"[{index_str}] Skip: {existing_files[0]}")
            continue

        print(f"[{index_str}] Downloading: {url} ({task})")
        try:
            ydl_opts = {
                'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
                'outtmpl': os.path.join(output_dir, f"{index_str}_{task}_%(title)s.%(ext)s"),
                'quiet': True,
                'no_warnings': True,
            }
            with YoutubeDL(ydl_opts) as ydl:
                ydl.download([url])
        except Exception as e:
            print(f"[{index_str}] Failed: {e}")

def build_dataset(url_file, output_root="."):
    video_dir = os.path.abspath(os.path.join(output_root, "video"))
    clip_dir = os.path.abspath(os.path.join(output_root, "video_clips"))
    
    urls = get_url_list(url_file)
    if not urls:
        print(f"Error: No valid data in {url_file}")
        return

    print(f"--- Step 1: Downloading {len(urls)} videos ---")
    download_videos(urls, video_dir)
    
    print(f"\n--- Step 2: Clipping videos ---")
    os.makedirs(clip_dir, exist_ok=True)
    for idx, item in enumerate(urls, 1):
        index_str = f"{idx:03d}"
        files = [f for f in os.listdir(video_dir) if f.startswith(f"{index_str}_")]
        if not files: continue
            
        input_file = os.path.join(video_dir, files[0])
        output_file = os.path.join(clip_dir, files[0])
        
        if os.path.exists(output_file): continue
            
        print(f"[{index_str}] Clipping: {files[0]}")
        center_sec = timestamp_to_seconds(item['timestamp'])
        clip_video(input_file, output_file, center_sec)

    print(f"\nDone! Clips saved in: {clip_dir}")

def main():
    import argparse
    parser = argparse.ArgumentParser(description='Build SBS Dataset from YouTube URL list')
    parser.add_argument('file', help='Path to youtube_url.txt')
    parser.add_argument('-o', '--output', default='.', help='Output root directory (default: .)')
    args = parser.parse_args()
    
    build_dataset(args.file, args.output)

if __name__ == "__main__":
    main()
