server.py

import os
import json
import glob
import subprocess
import shutil
import re
import argparse
import sys
import platform
import time
import urllib.request

# --- CONFIGURATIONS ---
TMP_DIR = "/tmp"
VRAM_DIR = os.path.expanduser("~/vram")
VRAM_SIZE_MB = 2500

KAGGLE_USER = os.getenv("KAGGLE_USERNAME")
KAGGLE_KEY = os.getenv("KAGGLE_KEY")

# v2.13.0: HF Space cookie refresh endpoint.
HF_SPACE_COOKIE_URL = os.getenv("HF_SPACE_COOKIE_URL", "https://sigyllly-test-final.hf.space/8w55da7/s591ad41a/request_net")


def log(task_id, msg):
    """Clean localized console logger formatting."""
    print(f"[{task_id}] {msg}", flush=True)


# ── Metadata Check & Smart Fallback Codecs ────────────────────────────────────

def check_video_metadata(task_id, url):
    """Asks yt-dlp to resolve the stream metadata IN ONE SHOT.
    Extracts the direct video/audio URLs and handles AV1->VP9 fallbacks directly,
    returning the exact stream dictionaries containing sizes and raw URLs.
    """
    log(task_id, "🔍 [1/6] Inspecting stream formats and extracting URLs (Single Shot)...")

    # Safe cookie pathing for the initial extraction
    cookies_path = os.path.join(os.getcwd(), "cookies.txt")
    cmd = [
        "python3", "-m", "yt_dlp",
        "-J",
        "-f", "bv*+ba/b",
        "--js-runtimes", "node",
        "--remote-components", "ejs:github",
        url
    ]
    if os.path.isfile(cookies_path):
        cmd.extend(["--cookies", cookies_path])

    process = subprocess.run(cmd, capture_output=True, text=True, check=True)
    meta = json.loads(process.stdout)
    formats = meta.get("formats", [])

    requested = meta.get("requested_formats")
    best_format = None
    a_info = None

    # Grab the resolved streams
    if requested:
        for f in requested:
            if f.get("vcodec") and f.get("vcodec") != "none":
                best_format = f
            elif f.get("acodec") and f.get("acodec") != "none":
                a_info = f

    # Single combined format fallback
    if best_format is None:
        if meta.get("vcodec") and meta.get("vcodec") != "none":
            best_format = meta

    # Fallback to manual audio format extraction if separated streams weren't grouped
    if a_info is None:
        audio_formats = [f for f in formats if f.get("vcodec") == "none" and f.get("acodec") != "none"]
        if audio_formats:
            audio_formats.sort(key=lambda x: x.get("abr", x.get("tbr", 0)), reverse=True)
            a_info = audio_formats[0]

    # Fallback to manual video format extraction if resolution failed
    if best_format is None:
        log(task_id, "⚠️ Could not read requested_formats from resolved selection; falling back to height/tbr sort.")
        video_formats = [f for f in formats if f.get("vcodec") != "none" and f.get("height")]
        if not video_formats:
            raise Exception("No explicit video streams discovered.")
        video_formats.sort(key=lambda x: (x.get("height", 0), x.get("tbr", 0)), reverse=True)
        best_format = video_formats[0]

    target_height = best_format.get("height", 0)
    vcodec = best_format.get("vcodec", "unknown")

    if "av01" in vcodec or "av1" in vcodec.lower():
        codec_label = "AV1"
    elif "avc" in vcodec or "264" in vcodec:
        codec_label = "H.264"
    elif "vp9" in vcodec or "vp09" in vcodec:
        codec_label = "VP9"
    else:
        codec_label = vcodec.split(".")[0].upper()

    filesize_bytes = best_format.get("filesize") or best_format.get("filesize_approx")
    estimated = False
    if not filesize_bytes:
        tbr = best_format.get("tbr")
        duration = meta.get("duration")
        if tbr and duration:
            filesize_bytes = tbr * 1000 / 8 * duration  # tbr is kbit/s
            estimated = True
            
    size_str = f"{filesize_bytes / (1024 * 1024):.1f} MB" + (" (est.)" if estimated else "") if filesize_bytes else "Unknown"

    log(task_id, f"📋 INITIAL ANALYSIS ➔ Best Format: {target_height}p | Codec: {codec_label} | Size: {size_str}")

    # AV1 to VP9 Fallback Processing
    if codec_label == "AV1":
        log(task_id, f"⚠️ AV1 detected at {target_height}p! Initiating search for VP9 alternative tracks...")
        video_formats = [f for f in formats if f.get("vcodec") != "none" and f.get("height")]
        vp9_formats = [
            f for f in video_formats
            if ("vp9" in f.get("vcodec", "").lower() or "vp09" in f.get("vcodec", "").lower())
        ]

        if not vp9_formats:
            raise ValueError("AV1 detected, and no alternative VP9 streams exist for this video.")

        exact_vp9 = [f for f in vp9_formats if f.get("height") == target_height]
        if exact_vp9:
            v_info = exact_vp9[0]
            log(task_id, f"🎯 Found exact VP9 match at matching resolution: {target_height}p")
        else:
            log(task_id, f"🔍 No exact {target_height}p VP9 track. Searching for nearest resolution >= 1080p...")
            eligible_vp9 = [f for f in vp9_formats if f.get("height") >= 1080]
            if not eligible_vp9:
                raise ValueError("AV1 fallback failed. Nearest available VP9 is lower than 1080p limit.")

            eligible_vp9.sort(key=lambda x: x.get("height", 0), reverse=True)
            v_info = eligible_vp9[0]

        final_height = v_info.get("height", 0)
        vp9_size_bytes = v_info.get("filesize") or v_info.get("filesize_approx")
        vp9_size_str = f"{vp9_size_bytes / (1024 * 1024):.1f} MB" if vp9_size_bytes else "Unknown"

        log(task_id, f"✅ SUCCESSFUL FALLBACK ROUTE ➔ Quality Tier: {final_height}p | Selected Codec: VP9 | Est. Size: {vp9_size_str}")
        print("-" * 75)
    else:
        v_info = best_format
        if target_height >= 2160: tier = "4K Ultra HD"
        elif target_height >= 1440: tier = "2K Quad HD"
        elif target_height >= 1080: tier = "1080p Full HD"
        elif target_height >= 720: tier = "720p HD"
        else: tier = f"{target_height}p SD"

        log(task_id, f"📋 SOURCE ANALYSIS ➔ Quality Tier: {tier} ({target_height}p) | Codec Profile: {codec_label} | Est. File Size: {size_str}")
        print("-" * 75)

    return v_info, a_info, meta


# ── Time & Setup Helpers for Aria2c ─────────────────────────────────────────

def time_to_seconds(t):
    if ':' in str(t):
        parts = str(t).split(':')
        if len(parts) == 3:
            return int(parts[0]) * 3600 + int(parts[1]) * 60 + float(parts[2])
        elif len(parts) == 2:
            return int(parts[0]) * 60 + float(parts[1])
    return float(t)


def get_actual_start_time(filename, stream_type="v:0"):
    cmd = [
        "ffprobe", "-v", "error", "-select_streams", stream_type, 
        "-show_entries", "packet=pts_time", "-of", "default=noprint_wrappers=1:nokey=1", 
        "-read_intervals", "%+#1", filename
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return float(result.stdout.strip())
    except Exception:
        return 0.0


def aria_download_cmd(cmd, task_id):
    """Infinite retry loop (Aria2c specific) with Smart Cookie Fallback."""
    CRITICAL_ERRORS = [2, 3, 4] 
    attempt = 1
    cookies_path = os.path.join(os.getcwd(), "cookies.txt")
    
    # Use a copy of the list so we can modify it dynamically during runtime
    current_cmd = list(cmd)
    
    while True:
        try:
            log(task_id, f"[Aria] Starting/Resuming chunk download...")
            subprocess.run(current_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
            log(task_id, "[✓] Download chunk successful.")
            break 
            
        except subprocess.CalledProcessError as e:
            # If it fails on the very first attempt and cookies are available, inject them for the retry
            if attempt == 1 and os.path.isfile(cookies_path) and "--load-cookies" not in current_cmd:
                log(task_id, f"[!] Download failed (Code {e.returncode}). Injecting cookies for next attempt...")
                current_cmd.extend(["--load-cookies", cookies_path])
                time.sleep(2)
            else:
                if e.returncode in CRITICAL_ERRORS:
                    log(task_id, f"[!] Critical Error {e.returncode}! System halting to prevent loop.")
                    raise e
                
                log(task_id, f"[!] Download interrupted (Code {e.returncode}). Resuming in 5 seconds...")
                time.sleep(5)
                
        attempt += 1


# ── Command Pipeline Stream Helpers ───────────────────────────────────────────

def parse_ffmpeg_progress(line, task_id):
    if "size=" in line or "time=" in line:
        clean_line = re.sub(r'\s+', ' ', line).strip()
        print(f"   ➔ [FFMPEG] {clean_line}", end="\r", flush=True)


def parse_kaggle_progress(line, task_id):
    if "%|" in line or "Upload successful" in line:
        print(f"   ➔ [KAG_UPLOAD] {line.strip()}", end="\r", flush=True)
    elif "datasets/joymant/" in line or "Dataset is being created" in line:
        print()
        log(task_id, line.strip())


def run_cmd(task_id, cmd, label):
    """Standard generic subprocess runner with progress trackers."""
    log(task_id, f">>> Starting {label}...")
    
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
    for line in process.stdout:
        line = line.rstrip()
        if line:
            if "ffmpeg" in label.lower():
                parse_ffmpeg_progress(line, task_id)
            elif "kaggle" in label.lower():
                parse_kaggle_progress(line, task_id)
            else:
                log(task_id, f"  {line}")

    process.wait()
    print()
    if process.returncode != 0:
        raise Exception(f"{label} failed with exit code {process.returncode}")
    log(task_id, f"<<< Finished {label} successfully.")


def upload_to_kaggle(task_id, file_path, filename):
    sanitized_id = re.sub(r'[^a-zA-Z0-9-]', '-', str(task_id)).lower().strip('-')
    upload_workspace = os.path.join(TMP_DIR, f"kaggle_upload_{sanitized_id}")
    os.makedirs(upload_workspace, exist_ok=True)
    shutil.move(file_path, os.path.join(upload_workspace, filename))

    with open(os.path.join(upload_workspace, "dataset-metadata.json"), "w") as f:
        json.dump({"title": f"Task Dataset {task_id}", "id": f"{KAGGLE_USER}/{sanitized_id}", "licenses": [{"name": "CC0-1.0"}]}, f, indent=4)

    run_cmd(task_id, ["kaggle", "datasets", "create", "-p", upload_workspace, "-r", "zip"], "Kaggle Dataset Upload")
    shutil.rmtree(upload_workspace, ignore_errors=True)


def force_cleanup():
    """Wipes Kaggle dirs, final vids, and aria temporary chunks"""
    try:
        subprocess.run(["sudo", "umount", "-f", VRAM_DIR], capture_output=True)
    except:
        pass
    if os.path.exists(VRAM_DIR):
        shutil.rmtree(VRAM_DIR, ignore_errors=True)
    for pattern in ["*.webm", "*.mkv", "*.mp4", "*.part", "*.ytdl", "*.tmp", "kaggle_upload_*"]:
        for f in glob.glob(os.path.join(TMP_DIR, pattern)):
            try:
                if os.path.isdir(f):
                    shutil.rmtree(f, ignore_errors=True)
                else:
                    os.remove(f)
            except:
                pass


# ── Core Task Process ─────────────────────────────────────────────────────────

def refresh_cookies_from_hf(task_id):
    """Call the HF Space to get fresh YouTube cookies via CloakBrowser."""
    log(task_id, "[Cookie Refresh] Requesting fresh cookies from HF Space...")
    try:
        req = urllib.request.Request(HF_SPACE_COOKIE_URL, method='GET')
        resp = urllib.request.urlopen(req, timeout=120)
        data = json.loads(resp.read().decode())
        if data.get('success') and data.get('cookies'):
            cookies_path = os.path.join(os.getcwd(), 'cookies.txt')
            with open(cookies_path, 'w') as f:
                f.write(data['cookies'])
            log(task_id, f"[Cookie Refresh] ✓ Fresh cookies saved to {cookies_path} ({len(data['cookies'])} bytes)")
            return True
        else:
            log(task_id, f"[Cookie Refresh] ✗ HF Space returned error: {data.get('error', 'unknown')}")
            return False
    except Exception as e:
        log(task_id, f"[Cookie Refresh] ✗ Failed to get cookies from HF Space: {e}")
        return False


def is_rate_limit_error(error_str):
    """Check if an error is a YouTube rate limit / 'Try again later' error.

    v2.14.0: Broadened detection — yt-dlp failures often just show as
    "returned non-zero exit status 1" without the actual HTTP error in
    the exception string. We now also catch generic yt-dlp command
    failures (which are almost always cookie/auth related) and trigger
    the cookie refresh.
    """
    lower = str(error_str).lower()
    # Direct rate limit markers
    rate_markers = ['429', 'too many requests', 'rate limit', 'try again',
                    'quota', 'http error 429', 'sign in to confirm', 'robot check',
                    'unable to extract', 'video unavailable', 'private video',
                    'members-only']
    if any(m in lower for m in rate_markers):
        return True
    # yt-dlp command failure (almost always cookie/auth related)
    if 'yt_dlp' in lower and 'non-zero exit status' in lower:
        return True
    if 'returned non-zero exit status' in lower and 'yt_dlp' in lower:
        return True
    return False


def process_task(task_id, url, start_time, end_time, filename, force_pipe):
    if not KAGGLE_USER or not KAGGLE_KEY:
        print("[CRITICAL ERROR] Kaggle variables missing!")
        return

    if not filename:
        filename = f"output_{task_id}.mkv"

    try:
        # 1. Fetch metadata and precise URLs in ONE shot
        v_info, a_info, info = check_video_metadata(task_id, url)
        
        if not v_info or not a_info:
            raise Exception("[!] Error: Separate video/audio streams not found for chunk downloading.")

        log(task_id, "⚡ Launching Hyper-Speed Aria2c Pipeline...")
        
        start_sec = time_to_seconds(start_time)
        end_sec = time_to_seconds(end_time)
        target_duration = end_sec - start_sec

        padding = 60
        padded_start = max(0, start_sec - padding)
        padded_end = end_sec + padding

        v_url, a_url = v_info['url'], a_info['url']
        
        v_size = v_info.get('filesize') or v_info.get('filesize_approx') or 0
        a_size = a_info.get('filesize') or a_info.get('filesize_approx') or 0
        v_dur = v_info.get('duration') or info.get('duration') or 1
        a_dur = a_info.get('duration') or info.get('duration') or 1

        # 2. Calculate Byte Ranges
        log(task_id, "[2/6] Calculating Bytes...")
        v_bps = v_size / v_dur if v_size else 0
        a_bps = a_size / a_dur if a_size else 0
        
        v_start_byte = int(padded_start * v_bps)
        v_end_byte = int(padded_end * v_bps)
        a_start_byte = int(padded_start * a_bps)
        a_end_byte = int(padded_end * a_bps)

        # 3. Mount Logic (VRAM vs TMP based on estimated fragment size instead of full file)
        chunk_size_mb = ((v_end_byte - v_start_byte) + (a_end_byte - a_start_byte)) / (1024 * 1024)
        use_pipe = force_pipe or (chunk_size_mb > (VRAM_SIZE_MB * 0.85))
        output_path = os.path.join(TMP_DIR if use_pipe else VRAM_DIR, filename)

        if not use_pipe:
            os.makedirs(VRAM_DIR, exist_ok=True)
            subprocess.run(["sudo", "mount", "-t", "tmpfs", "-o", f"size={VRAM_SIZE_MB}M", "tmpfs", VRAM_DIR])

        # 4. Initiate 16-Connection Download
        log(task_id, f"[3/6] Hyper-Speed Download to {TMP_DIR} (16 Connections, Auto-Resume)...")
        
        # Base Aria config (no cookies loaded here initially)
        base_aria = ["aria2c", "-c", "-x", "16", "-s", "16", "--file-allocation=none", "-d", TMP_DIR]
        
        aria_download_cmd(base_aria + ["-o", "v_init.tmp", f"{v_url}&range=0-5000"], task_id)
        aria_download_cmd(base_aria + ["-o", "v_chunk.tmp", f"{v_url}&range={v_start_byte}-{v_end_byte}"], task_id)
        aria_download_cmd(base_aria + ["-o", "a_init.tmp", f"{a_url}&range=0-5000"], task_id)
        aria_download_cmd(base_aria + ["-o", "a_chunk.tmp", f"{a_url}&range={a_start_byte}-{a_end_byte}"], task_id)

        # 5. Concatenate Init Blocks with Body Blocks
        log(task_id, "[4/6] Internal CAT Merge in /tmp...")
        v_rough = os.path.join(TMP_DIR, "rough_v.tmp")
        a_rough = os.path.join(TMP_DIR, "rough_a.tmp")
        
        with open(v_rough, 'wb') as outfile:
            for f in ["v_init.tmp", "v_chunk.tmp"]:
                with open(os.path.join(TMP_DIR, f), 'rb') as infile: outfile.write(infile.read())
                
        with open(a_rough, 'wb') as outfile:
            for f in ["a_init.tmp", "a_chunk.tmp"]:
                with open(os.path.join(TMP_DIR, f), 'rb') as infile: outfile.write(infile.read())

        # 6. Retrieve proper offsets for precise sync
        log(task_id, "[5/6] Probing actual start times...")
        v_actual_start = get_actual_start_time(v_rough, "v:0")
        a_actual_start = get_actual_start_time(a_rough, "a:0")
        
        v_offset = max(0, start_sec - v_actual_start)
        a_offset = max(0, start_sec - a_actual_start)
        log(task_id, f"      Video Offset: {v_offset:.3f}s | Audio Offset: {a_offset:.3f}s")

        # 7. Final FFMPEG Cut & Copy
        log(task_id, f"[6/6] Final FFmpeg Precision Trim ➔ {output_path}")
        ffmpeg_path = shutil.which("ffmpeg") or os.path.expanduser("~/.local/bin/ffmpeg")
        
        ffmpeg_cmd = [
            ffmpeg_path, "-y", 
            "-loglevel", "error", "-stats",
            "-ss", str(v_offset), "-i", v_rough,
            "-ss", str(a_offset), "-i", a_rough,
            "-t", str(target_duration),
            "-c", "copy",
            output_path
        ]
        run_cmd(task_id, ffmpeg_cmd, "ffmpeg clip split")

        log(task_id, "Transporting file payload over to Kaggle...")
        upload_to_kaggle(task_id, output_path, filename)
        log(task_id, "🏁 SUCCESS: Task completed.")

    except Exception as e:
        err_str = str(e)
        if is_rate_limit_error(err_str):
            log(task_id, f"⚠️ Rate limit detected. Attempting cookie refresh from HF Space...")
            if refresh_cookies_from_hf(task_id):
                log(task_id, "🔄 Retrying download with fresh cookies...")
                try:
                    # Retry the full pipeline once with fresh cookies
                    v_info, a_info, info = check_video_metadata(task_id, url)
                    if not v_info or not a_info:
                        raise Exception("Separate video/audio streams not found.")

                    log(task_id, "⚡ Retrying Hyper-Speed Aria2c Pipeline...")
                    start_sec = time_to_seconds(start_time)
                    end_sec = time_to_seconds(end_time)
                    target_duration = end_sec - start_sec
                    padding = 60
                    padded_start = max(0, start_sec - padding)
                    padded_end = end_sec + padding
                    v_url, a_url = v_info['url'], a_info['url']
                    v_size = v_info.get('filesize') or v_info.get('filesize_approx') or 0
                    a_size = a_info.get('filesize') or a_info.get('filesize_approx') or 0
                    v_dur = v_info.get('duration') or info.get('duration') or 1
                    a_dur = a_info.get('duration') or info.get('duration') or 1
                    v_bps = v_size / v_dur if v_size else 0
                    a_bps = a_size / a_dur if a_size else 0
                    v_start_byte = int(padded_start * v_bps)
                    v_end_byte = int(padded_end * v_bps)
                    a_start_byte = int(padded_start * a_bps)
                    a_end_byte = int(padded_end * a_bps)
                    chunk_size_mb = ((v_end_byte - v_start_byte) + (a_end_byte - a_start_byte)) / (1024 * 1024)
                    use_pipe = force_pipe or (chunk_size_mb > (VRAM_SIZE_MB * 0.85))
                    output_path = os.path.join(TMP_DIR if use_pipe else VRAM_DIR, filename)
                    if not use_pipe:
                        os.makedirs(VRAM_DIR, exist_ok=True)
                        subprocess.run(["sudo", "mount", "-t", "tmpfs", "-o", f"size={VRAM_SIZE_MB}M", "tmpfs", VRAM_DIR])

                    cookies_path = os.path.join(os.getcwd(), "cookies.txt")
                    base_aria = ["aria2c", "-c", "-x", "16", "-s", "16", "--file-allocation=none", "-d", TMP_DIR]
                    if os.path.isfile(cookies_path):
                        base_aria.extend(["--load-cookies", cookies_path])

                    aria_download_cmd(base_aria + ["-o", "v_init.tmp", f"{v_url}&range=0-5000"], task_id)
                    aria_download_cmd(base_aria + ["-o", "v_chunk.tmp", f"{v_url}&range={v_start_byte}-{v_end_byte}"], task_id)
                    aria_download_cmd(base_aria + ["-o", "a_init.tmp", f"{a_url}&range=0-5000"], task_id)
                    aria_download_cmd(base_aria + ["-o", "a_chunk.tmp", f"{a_url}&range={a_start_byte}-{a_end_byte}"], task_id)

                    v_rough = os.path.join(TMP_DIR, "rough_v.tmp")
                    a_rough = os.path.join(TMP_DIR, "rough_a.tmp")
                    with open(v_rough, 'wb') as outfile:
                        for f in ["v_init.tmp", "v_chunk.tmp"]:
                            with open(os.path.join(TMP_DIR, f), 'rb') as infile: outfile.write(infile.read())
                    with open(a_rough, 'wb') as outfile:
                        for f in ["a_init.tmp", "a_chunk.tmp"]:
                            with open(os.path.join(TMP_DIR, f), 'rb') as infile: outfile.write(infile.read())

                    v_actual_start = get_actual_start_time(v_rough, "v:0")
                    a_actual_start = get_actual_start_time(a_rough, "a:0")
                    v_offset = max(0, start_sec - v_actual_start)
                    a_offset = max(0, start_sec - a_actual_start)

                    ffmpeg_path = shutil.which("ffmpeg") or os.path.expanduser("~/.local/bin/ffmpeg")
                    ffmpeg_cmd = [ffmpeg_path, "-y", "-loglevel", "error", "-stats",
                                  "-ss", str(v_offset), "-i", v_rough,
                                  "-ss", str(a_offset), "-i", a_rough,
                                  "-t", str(target_duration), "-c", "copy", output_path]
                    run_cmd(task_id, ffmpeg_cmd, "ffmpeg clip split (retry)")
                    upload_to_kaggle(task_id, output_path, filename)
                    log(task_id, "🏁 SUCCESS: Task completed (with refreshed cookies).")
                except Exception as retry_e:
                    log(task_id, f"❌ PIPELINE FAILED: Retry with fresh cookies also failed. Try after later. Error: {retry_e}")
            else:
                log(task_id, f"❌ PIPELINE FAILED: Cookie refresh failed. Try after later. Error: {e}")
        else:
            log(task_id, f"❌ PIPELINE FAILED: Try again 1 hour later. Error: {e}")
    finally:
        force_cleanup()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Local Video Export Tool")
    parser.add_argument("--id", required=True)
    parser.add_argument("--url", required=True)
    parser.add_argument("--start", required=True)
    parser.add_argument("--end", required=True)
    parser.add_argument("--name", default=None)
    parser.add_argument("--force-pipe", action="store_true")

    args = parser.parse_args()
    process_task(args.id, args.url, args.start, args.end, args.name, args.force_pipe)