#!/usr/bin/env python3
"""Download multiple YouTube videos via vidssave API in parallel. Always 256KBPS."""

import os, json, time, requests, re, sys
from concurrent.futures import ThreadPoolExecutor, as_completed

AUTH = "20250901majwlqo"
DOMAIN = "api-ak.vidssave.com"
BASE_DIR = "/home/ubuntu/Speech_maker_pipeline/pawan_kalyan"

URLS = [
    "https://www.youtube.com/watch?v=gr7TttotBD8",
    "https://www.youtube.com/watch?v=VSyivphv2J4",
    "https://www.youtube.com/watch?v=WgcNsi6VFtY",
    "https://www.youtube.com/watch?v=42AMTLWpZ9A",
    "https://www.youtube.com/watch?v=WJSZ5OuIYkc",
    "https://www.youtube.com/watch?v=uPdXi6xIBlc",
    "https://www.youtube.com/watch?v=p9xrgKSt82o",
    "https://www.youtube.com/watch?v=zAtgTpO-724",
    "https://www.youtube.com/watch?v=68wSTlrw_RQ",
    "https://www.youtube.com/watch?v=WdJQGbeBdyw",
    "https://www.youtube.com/watch?v=R8E5ihYT_KU",
    "https://www.youtube.com/watch?v=lcWLYBj7Jhc",
    "https://www.youtube.com/watch?v=HmqZfwSIOmU",
    "https://www.youtube.com/watch?v=zome0ret9zA",
    "https://www.youtube.com/watch?v=rOiJvev-2LU",
]

def extract_video_id(url):
    m = re.search(r'[?&]v=([^&]+)', url)
    return m.group(1) if m else url.split('/')[-1]

def download_video(url):
    vid = extract_video_id(url)
    out_dir = os.path.join(BASE_DIR, vid)
    os.makedirs(out_dir, exist_ok=True)

    # Check if already downloaded
    existing = [f for f in os.listdir(out_dir) if f.startswith("yt_downloaded_256kbps")]
    if existing:
        print(f"[{vid}] Already downloaded, skipping")
        return vid, True, "already exists"

    # Step 1: Parse
    try:
        resp = requests.post("https://api.vidssave.com/api/contentsite_api/media/parse", data={
            "auth": AUTH, "domain": DOMAIN, "origin": "source", "link": url
        }, timeout=30)
        data = resp.json()
    except Exception as e:
        return vid, False, f"parse error: {e}"

    if data.get("status") != 1:
        return vid, False, f"parse failed: {data}"

    title = data["data"].get("title", "unknown")
    duration = data["data"].get("duration", 0)
    resources = data["data"].get("resources", [])

    # Find 256KBPS audio
    target = None
    for r in resources:
        if r["type"] == "audio" and r["quality"] == "256KBPS":
            target = r
            break
    # Fallback to 128KBPS if 256 not available
    if not target:
        for r in resources:
            if r["type"] == "audio" and r["quality"] == "128KBPS":
                target = r
                break
    # Fallback to any audio
    if not target:
        for r in resources:
            if r["type"] == "audio":
                target = r
                break

    if not target:
        return vid, False, "no audio resource found"

    quality = target["quality"]
    fmt = target["format"].lower()
    resource_content = target["resource_content"]

    # If direct download available
    if target.get("download_url"):
        dl_url = target["download_url"]
    else:
        # Step 2: Request download
        try:
            resp = requests.post("https://api.vidssave.com/api/contentsite_api/media/download", data={
                "auth": AUTH, "domain": DOMAIN,
                "request": resource_content, "no_encrypt": "1"
            }, timeout=30)
            dl_data = resp.json()
        except Exception as e:
            return vid, False, f"download request error: {e}"

        if dl_data.get("status") != 1:
            return vid, False, f"download request failed: {dl_data}"

        task_id = dl_data["data"]["task_id"]

        # Step 3: Poll SSE for download link
        sse_url = (
            f"https://api.vidssave.com/sse/contentsite_api/media/download_query"
            f"?auth={AUTH}&domain={DOMAIN}&task_id={task_id}"
            f"&download_domain=vidssave.com&origin=content_site"
        )
        dl_url = None
        try:
            resp = requests.get(sse_url, stream=True, timeout=120)
            for line in resp.iter_lines(decode_unicode=True):
                if not line:
                    continue
                if line.startswith("data:"):
                    event_data = json.loads(line[5:].strip())
                    if event_data.get("status") == "success":
                        dl_url = event_data.get("download_link")
                        break
                    elif event_data.get("status") == "error":
                        return vid, False, f"SSE error: {event_data}"
        except Exception as e:
            return vid, False, f"SSE poll error: {e}"

        if not dl_url:
            return vid, False, "no download link from SSE"

    # Step 4: Download file
    ext_map = {"opus": "opus", "m4a": "m4a", "webm": "webm", "mp4": "m4a"}
    ext = ext_map.get(fmt, fmt)
    out_path = os.path.join(out_dir, f"yt_downloaded_256kbps.{ext}")

    try:
        resp = requests.get(dl_url, stream=True, timeout=300, allow_redirects=True)
        resp.raise_for_status()
        with open(out_path, "wb") as f:
            for chunk in resp.iter_content(chunk_size=1024*1024):
                f.write(chunk)
    except Exception as e:
        return vid, False, f"download error: {e}"

    size_mb = os.path.getsize(out_path) / (1024*1024)
    print(f"[{vid}] OK - {quality} {fmt} - {size_mb:.1f}MB - {duration}s - {title[:50]}")
    return vid, True, f"{size_mb:.1f}MB"

if __name__ == "__main__":
    print(f"Downloading {len(URLS)} videos in parallel...\n")
    
    results = {}
    with ThreadPoolExecutor(max_workers=5) as pool:
        futures = {pool.submit(download_video, url): url for url in URLS}
        for future in as_completed(futures):
            vid, ok, msg = future.result()
            results[vid] = (ok, msg)
            if not ok:
                print(f"[{vid}] FAILED: {msg}")

    print(f"\n{'='*60}")
    ok_count = sum(1 for v in results.values() if v[0])
    print(f"Done: {ok_count}/{len(URLS)} succeeded")
    for vid, (ok, msg) in sorted(results.items()):
        status = "OK" if ok else "FAIL"
        print(f"  [{status}] {vid}: {msg}")