#!/usr/bin/env python3
"""
YT Video Analysis Worker v2 — Runs on Kaggle None runtime
Constraints: 30GB RAM · 20GB disk · 4 vCPUs · CPU-only
Pipeline: Download → Split Video 60s → Compress → Upload chunks to HF bucket
          → Vision model per chunk → Extract audio → Split if >3min
          → Transcribe → LLM analysis → Collect → Parquet → Upload dataset
"""

import base64, glob, hashlib, json, logging, os, re, shutil, subprocess
import sys, tempfile, time, traceback, uuid
from datetime import datetime, timezone
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("worker")

# ── constants ────────────────────────────────────────────────────────
HF_REPO          = "AdhyanshVerma/YT"
STAGING_PREFIX   = "staging"
DATA_PREFIX      = "data"
VIDEO_CHUNK_SEC  = 60
AUDIO_SPLIT_SEC  = 180       # 3 minutes
MIN_TAIL_SEC     = 15
FRAME_INTERVAL   = 5         # extract 1 frame per N seconds
MAX_RETRIES      = 3
RETRY_DELAY      = 5
DISK_HEADROOM_GB = 4

# ── disk helpers ─────────────────────────────────────────────────────
def disk_free_gb(path="/"):
    st = os.statvfs(path)
    return (st.f_bavail * st.f_frsize) / (1024**3)

def safe_cleanup(*paths):
    for p in paths:
        try:
            p = str(p)
            if os.path.isdir(p):
                shutil.rmtree(p, ignore_errors=True)
            elif os.path.isfile(p):
                os.remove(p)
        except Exception:
            pass

# ── ffprobe duration ─────────────────────────────────────────────────
def media_duration(path):
    try:
        r = subprocess.run(
            ["ffprobe","-v","error","-show_entries","format=duration",
             "-of","default=noprint_wrappers=1:nokey=1", path],
            capture_output=True, text=True, check=True)
        return float(r.stdout.strip())
    except Exception:
        return None

# ── video download ───────────────────────────────────────────────────
def download_video(url, out_dir):
    """Download video at HD, return filepath or None."""
    os.makedirs(out_dir, exist_ok=True)
    
    # Path to cookies file (useful for bypassing bot blocks)
    # Checks either current working dir or the parent of the script
    cookies_path = "cookies.txt"
    if not os.path.exists(cookies_path):
        # Safely get directory, falling back if __file__ is not defined (e.g. in notebooks)
        try:
            base_d = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            alt_cookies = os.path.join(base_d, "cookies.txt")
            if os.path.exists(alt_cookies):
                cookies_path = alt_cookies
        except NameError:
            pass

    opts = {
        "format": "bestvideo[height<=1080]+bestaudio/best[height<=1080]/bestvideo+bestaudio/best",
        "outtmpl": os.path.join(out_dir, "%(id)s.%(ext)s"),
        "merge_output_format": "mp4",
        "quiet": True, "no_warnings": True,
        "nocheckcertificate": True, "ignoreerrors": False,
        "geo_bypass": True,
        "socket_timeout": 60,
        "retries": 5,
        "sleep_requests": 2,
        "sleep_interval": 5,
        "max_sleep_interval": 15,
        "extractor_args": {"youtube": {"player_client": ["android", "ios", "tv", "web"]}},
        "writeinfojson": True,
        "writecomments": True,
        "writesubtitles": True,
        "subtitleslangs": ["en", "auto"],
        "http_headers": {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                          "AppleWebKit/537.36 Chrome/125.0 Safari/537.36"
        },
    }
    
    # if os.path.exists(cookies_path):
    #     opts["cookiefile"] = cookies_path
        
    try:
        import yt_dlp
        with yt_dlp.YoutubeDL(opts) as ydl:
            info = ydl.extract_info(url, download=True)
            if not info:
                return None, {}
            fp = None
            if "requested_downloads" in info:
                fp = info["requested_downloads"][0]["filepath"]
            else:
                fp = ydl.prepare_filename(info)
                fp = os.path.splitext(fp)[0] + ".mp4"
            meta = {
                "title": info.get("title",""),
                "duration": info.get("duration",0),
                "duration_string": info.get("duration_string",""),
                "view_count": info.get("view_count",0),
                "like_count": info.get("like_count",0),
                "comment_count": info.get("comment_count",0),
                "channel": info.get("channel",""),
                "channel_id": info.get("channel_id",""),
                "channel_url": info.get("channel_url",""),
                "channel_follower_count": info.get("channel_follower_count",0),
                "upload_date": info.get("upload_date",""),
                "timestamp": info.get("timestamp",0),
                "description": info.get("description",""),
                "tags": info.get("tags",[]),
                "categories": info.get("categories",[]),
                "language": info.get("language",""),
                "license": info.get("license",""),
                "availability": info.get("availability",""),
                "age_limit": info.get("age_limit",0),
                "width": info.get("width",0),
                "height": info.get("height",0),
                "resolution": info.get("resolution",""),
                "fps": info.get("fps",0),
                "aspect_ratio": info.get("aspect_ratio",0),
                "vcodec": info.get("vcodec",""),
                "acodec": info.get("acodec",""),
                "vbr": info.get("vbr",0),
                "abr": info.get("abr",0),
                "asr": info.get("asr",0),
                "audio_channels": info.get("audio_channels",0),
                "dynamic_range": info.get("dynamic_range",""),
                "filesize_approx": info.get("filesize_approx",0),
                "tbr": info.get("tbr",0),
                "ext": info.get("ext",""),
                "protocol": info.get("protocol",""),
                "format_id": info.get("format_id",""),
                "format_note": info.get("format_note",""),
                "playlist_title": info.get("playlist_title",""),
                "playlist_index": info.get("playlist_index",0),
                "live_status": info.get("live_status",""),
                "video_id": info.get("id",""),
                "webpage_url": info.get("webpage_url", url),
            }
            return fp, meta
    except Exception as e:
        log.error(f"Download failed for {url}: {e}")
        return None, {}

# ── video splitting ──────────────────────────────────────────────────
def split_video(input_path, chunk_sec=VIDEO_CHUNK_SEC, out_dir=None):
    name = Path(input_path).stem
    if out_dir is None:
        out_dir = os.path.join(os.path.dirname(input_path), f"{name}_vchunks")
    os.makedirs(out_dir, exist_ok=True)
    dur = media_duration(input_path)
    if not dur:
        return []
    # calculate split points
    pts, cur = [], chunk_sec
    while cur < dur:
        if (dur - cur) < MIN_TAIL_SEC:
            break
        pts.append(cur); cur += chunk_sec
    pattern = os.path.join(out_dir, f"{name}_chunk_%03d.mp4")
    cmd = ["ffmpeg","-y","-i",input_path,"-map","0","-c","copy",
           "-f","segment","-reset_timestamps","1"]
    if pts:
        cmd += ["-segment_times", ",".join(str(round(p,3)) for p in pts)]
    else:
        cmd += ["-segment_time", str(dur+10)]
    cmd.append(pattern)
    try:
        subprocess.run(cmd, check=True, capture_output=True)
    except subprocess.CalledProcessError as e:
        log.error(f"Video split error: {e.stderr.decode(errors='replace')[:300]}")
        return []
    return sorted(glob.glob(os.path.join(out_dir, f"{name}_chunk_*.mp4")))

# ── video compression (same resolution, smaller size) ────────────────
def compress_chunk(inp, out):
    """Re-encode with h264 CRF 28 — same res, ~40-60% smaller."""
    cmd = ["ffmpeg","-y","-i",inp,
           "-c:v","libx264","-crf","28","-preset","fast",
           "-c:a","aac","-b:a","128k", out]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        return True
    except Exception:
        # if compression fails, just copy
        shutil.copy2(inp, out)
        return False

# ── audio extraction ─────────────────────────────────────────────────
def extract_audio(video_path, out_path):
    cmd = ["ffmpeg","-y","-i",video_path,
           "-vn","-acodec","libmp3lame","-q:a","4", out_path]
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        return True
    except subprocess.CalledProcessError as e:
        log.error(f"Audio extraction failed for {video_path}: {e.stderr}")
        return False
    except Exception as e:
        log.error(f"Unexpected error in audio extraction: {e}")
        return False

# ── audio silence detection + splitting ──────────────────────────────
def detect_silences(audio_path):
    cmd = ["ffmpeg","-i",audio_path,
           "-af","silencedetect=noise=-20dB:d=0.3","-f","null","-"]
    r = subprocess.run(cmd, capture_output=True, text=True)
    starts, ends = [], []
    for line in r.stderr.split("\n"):
        m = re.search(r"silence_start:\s+([\d.]+)", line)
        if m: starts.append(float(m.group(1)))
        m = re.search(r"silence_end:\s+([\d.]+)", line)
        if m: ends.append(float(m.group(1)))
    return [(s+e)/2 for s,e in zip(starts, ends)]

def split_audio(audio_path, chunk_sec=AUDIO_SPLIT_SEC, out_dir=None):
    name = Path(audio_path).stem
    if out_dir is None:
        out_dir = os.path.join(os.path.dirname(audio_path), f"{name}_achunks")
    os.makedirs(out_dir, exist_ok=True)
    dur = media_duration(audio_path)
    if not dur or dur <= chunk_sec:
        return [audio_path]  # no split needed
    silences = detect_silences(audio_path)
    pts, target = [], chunk_sec
    while target < dur:
        if (dur - target) < MIN_TAIL_SEC:
            break
        valid = [p for p in silences if abs(p - target) <= 15]
        if valid:
            best = min(valid, key=lambda p: abs(p - target))
            pts.append(best); target = best + chunk_sec
        else:
            pts.append(target); target += chunk_sec
    ext = Path(audio_path).suffix
    pattern = os.path.join(out_dir, f"{name}_achunk_%03d{ext}")
    cmd = ["ffmpeg","-y","-i",audio_path,"-c","copy","-map","0",
           "-f","segment","-reset_timestamps","1"]
    if pts:
        cmd += ["-segment_times", ",".join(str(round(p,3)) for p in pts)]
    else:
        cmd += ["-segment_time", str(dur+10)]
    cmd.append(pattern)
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        return sorted(glob.glob(os.path.join(out_dir, f"{name}_achunk_*{ext}")))
    except Exception:
        return [audio_path]

# ── frame extraction for vision API ──────────────────────────────────
def extract_frames(video_path, out_dir, interval=FRAME_INTERVAL):
    os.makedirs(out_dir, exist_ok=True)
    pattern = os.path.join(out_dir, "frame_%04d.jpg")
    cmd = ["ffmpeg","-y","-i",video_path,"-vf",f"fps=1/{interval}",
           "-q:v","5", pattern]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        return sorted(glob.glob(os.path.join(out_dir, "frame_*.jpg")))
    except Exception:
        return []

def encode_image_b64(path):
    with open(path, "rb") as f:
        return "data:image/jpeg;base64," + base64.b64encode(f.read()).decode()

# ── HF storage bucket ───────────────────────────────────────────────
def hf_upload_file(local_path, remote_path, token):
    from huggingface_hub import HfApi
    api = HfApi(token=token)
    for attempt in range(MAX_RETRIES):
        try:
            api.upload_file(
                path_or_fileobj=local_path,
                path_in_repo=remote_path,
                repo_id=HF_REPO, repo_type="dataset",
            )
            return True
        except Exception as e:
            log.warning(f"HF upload attempt {attempt+1} failed: {e}")
            time.sleep(RETRY_DELAY * (attempt+1))
    return False

def hf_download_url(remote_path):
    return f"https://huggingface.co/datasets/{HF_REPO}/resolve/main/{remote_path}"

def hf_delete_folder(prefix, token):
    from huggingface_hub import HfApi
    api = HfApi(token=token)
    try:
        files = api.list_repo_files(repo_id=HF_REPO, repo_type="dataset")
        to_del = [f for f in files if f.startswith(prefix)]
        if to_del:
            ops = []
            from huggingface_hub import CommitOperationDelete
            for f in to_del:
                ops.append(CommitOperationDelete(path_in_repo=f))
            api.create_commit(
                repo_id=HF_REPO, repo_type="dataset",
                operations=ops, commit_message=f"cleanup {prefix}")
    except Exception as e:
        log.warning(f"HF cleanup failed: {e}")

# ── model API calls ──────────────────────────────────────────────────
def _api_call(messages, api_key, model, base_url="https://api.featherless.ai/v1",
              max_tokens=4096, temperature=0.4):
    """OpenAI-compatible API call with retries."""
    from openai import OpenAI
    client = OpenAI(base_url=base_url, api_key=api_key)
    for attempt in range(MAX_RETRIES):
        try:
            resp = client.chat.completions.create(
                model=model, messages=messages,
                max_tokens=max_tokens, temperature=temperature)
            return resp.choices[0].message.content
        except Exception as e:
            log.warning(f"API attempt {attempt+1} failed ({model}): {e}")
            time.sleep(RETRY_DELAY * (attempt+1))
    return None

DEFAULT_VISION_MODELS = [
    "MiniMaxAI/MiniMax-M3",
    "XiaomiMiMo/MiMo-V2.5",
    "Qwen/Qwen3-VL-32B-Instruct",
    "Qwen/Qwen3.6-27B",
    "Qwen/Qwen3-VL-235B-A22B-Thinking",
]

def analyze_video_chunk(frames_dir, chunk_idx, total_chunks, api_key,
                        vision_models=None):
    """Send extracted frames to vision model for analysis with priority fallback."""
    frames = sorted(glob.glob(os.path.join(frames_dir, "frame_*.jpg")))
    if not frames:
        return None
    prompt = (
        f"You are analyzing video chunk {chunk_idx+1}/{total_chunks}. "
        f"These are frames extracted every {FRAME_INTERVAL}s from a 60-second segment.\n\n"
        "Provide an EXTREMELY detailed analysis:\n"
        "1. Scene description — what is happening visually\n"
        "2. Objects, people, text/UI visible on screen\n"
        "3. Motion and transitions between frames\n"
        "4. Any on-screen text, code, diagrams, or slides\n"
        "5. Visual style, colors, camera angles\n"
        "6. Key information conveyed visually\n\n"
        "Be thorough — every detail matters for the dataset."
    )
    content = [{"type": "text", "text": prompt}]
    # send up to 12 frames to avoid token limits
    step = max(1, len(frames) // 12)
    selected = frames[::step][:12]
    for fp in selected:
        content.append({
            "type": "image_url",
            "image_url": {"url": encode_image_b64(fp)}
        })
    messages = [{"role": "user", "content": content}]

    models = vision_models if vision_models else DEFAULT_VISION_MODELS
    for model in models:
        log.info(f"Attempting vision analysis with model: {model}")
        res = _api_call(messages, api_key, model, max_tokens=4096)
        if res:
            log.info(f"Vision analysis succeeded with model: {model}")
            return res
        log.warning(f"Vision model {model} failed or busy. Trying next in priority...")
    return None


def analyze_audio_transcript(transcript, chunk_idx, total_chunks, api_key,
                             model="Qwen/Qwen3-32B"):
    """Send audio transcript to LLM for detailed analysis."""
    prompt = (
        f"Audio transcript chunk {chunk_idx+1}/{total_chunks}:\n\n"
        f"---\n{transcript}\n---\n\n"
        "Provide a detailed analysis of this audio:\n"
        "1. Summary of what is being discussed\n"
        "2. Key points and arguments made\n"
        "3. Technical terms or concepts mentioned\n"
        "4. Speaker tone and intent\n"
        "5. Notable quotes or statements\n"
        "6. Questions raised or answered\n\n"
        "Be comprehensive and precise."
    )
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    return _api_call(messages, api_key, model, max_tokens=4096)

def transcribe_audio_chunk(audio_path):
    """Transcribe audio using faster-whisper (CPU)."""
    try:
        from faster_whisper import WhisperModel
        model = WhisperModel("base", device="cpu", compute_type="int8")
        segments, info = model.transcribe(audio_path, beam_size=3)
        text_parts = []
        for seg in segments:
            text_parts.append(f"[{seg.start:.1f}s-{seg.end:.1f}s] {seg.text.strip()}")
        return "\n".join(text_parts)
    except Exception as e:
        log.error(f"Transcription failed: {e}")
        return None

# ── per-video pipeline ───────────────────────────────────────────────
def process_video(video_url, video_id, api_key, hf_token, worker_id, work_dir,
                  vision_models=None):
    """Full pipeline for one video. Returns result dict or None."""
    log.info(f"{'='*60}")
    log.info(f"Processing: {video_id} — {video_url}")
    log.info(f"Disk free: {disk_free_gb():.1f} GB")

    vid_dir = os.path.join(work_dir, video_id)
    os.makedirs(vid_dir, exist_ok=True)
    staging_prefix = f"{STAGING_PREFIX}/{worker_id}/{video_id}"
    
    result = {
        "video_id": video_id,
        "url": video_url,
        "worker_id": worker_id,
        "processed_at": datetime.now(timezone.utc).isoformat(),
        "status": "failed",
        "video_analysis": [],
        "audio_analysis": [],
        "audio_transcript": "",
        "metadata": {},
        "hf_video_urls": [],
        "errors": [],
    }

    try:
        # ── 1. Download ──────────────────────────────────────────
        log.info("Step 1: Downloading video...")
        dl_dir = os.path.join(vid_dir, "download")
        video_path, meta = download_video(video_url, dl_dir)
        if not video_path or not os.path.exists(video_path):
            log.error(f"Download failed for {video_id}. Skipping commit to avoid empty records.")
            return None
        result["metadata"] = meta
        fsize = os.path.getsize(video_path) / (1024**2)
        log.info(f"Downloaded: {fsize:.0f} MB — {meta.get('title','?')}")

        # ── 2. Split video into 60s chunks (skip for short videos) ──
        vid_duration = media_duration(video_path) or meta.get("duration", 0)
        if vid_duration and vid_duration <= VIDEO_CHUNK_SEC + MIN_TAIL_SEC:
            log.info(f"Step 2: Short video ({vid_duration:.0f}s ≤ {VIDEO_CHUNK_SEC}s) — skipping split")
            video_chunks = [video_path]
        else:
            log.info(f"Step 2: Splitting {vid_duration:.0f}s video into {VIDEO_CHUNK_SEC}s chunks...")
            chunks_dir = os.path.join(vid_dir, "vchunks")
            video_chunks = split_video(video_path, VIDEO_CHUNK_SEC, chunks_dir)
            if not video_chunks:
                video_chunks = [video_path]
        log.info(f"Video chunks: {len(video_chunks)}")

        # ── 3. Compress + upload chunks + vision analysis ────────
        log.info("Step 3: Compress → Upload → Vision analysis...")
        for cidx, chunk_path in enumerate(video_chunks):
            log.info(f"  Chunk {cidx+1}/{len(video_chunks)}: {Path(chunk_path).name}")
            # compress
            comp_path = chunk_path.replace(".mp4", "_comp.mp4")
            compress_chunk(chunk_path, comp_path)
            if os.path.exists(comp_path):
                if chunk_path != video_path:
                    safe_cleanup(chunk_path)
                chunk_path = comp_path

            # upload to HF bucket
            remote = f"{staging_prefix}/v_chunk_{cidx:03d}.mp4"
            uploaded = hf_upload_file(chunk_path, remote, hf_token)
            if not uploaded:
                result["errors"].append(f"Upload failed chunk {cidx}")
                if chunk_path != video_path:
                    safe_cleanup(chunk_path)
                continue

            # extract frames for vision API
            result["hf_video_urls"].append(hf_download_url(remote))
            frames_dir = os.path.join(vid_dir, f"frames_{cidx}")
            frames = extract_frames(chunk_path, frames_dir)
            if chunk_path != video_path:
                safe_cleanup(chunk_path)  # free disk

            if frames:
                analysis = analyze_video_chunk(
                    frames_dir, cidx, len(video_chunks), api_key, vision_models)
                result["video_analysis"].append({
                    "chunk_index": cidx,
                    "chunk_url": hf_download_url(remote),
                    "analysis": analysis or "ANALYSIS_FAILED",

                    "frame_count": len(frames),
                })
            safe_cleanup(frames_dir)

        # ── 4. Extract audio ─────────────────────────────────────
        log.info("Step 4: Extracting audio...")
        audio_path = os.path.join(vid_dir, f"{video_id}.mp3")
        audio_ok = extract_audio(video_path, audio_path)
        # We no longer delete the original video here so it can be stored in the scalable directory system
        # safe_cleanup(video_path, dl_dir)

        if audio_ok and os.path.exists(audio_path):
            audio_dur = media_duration(audio_path)
            log.info(f"Audio duration: {audio_dur:.0f}s")

            # ── 5. Split audio if > 3 min ────────────────────────
            if audio_dur and audio_dur > AUDIO_SPLIT_SEC:
                log.info("Step 5: Splitting audio (>3min)...")
                audio_chunks = split_audio(audio_path, AUDIO_SPLIT_SEC)
                if audio_chunks and audio_chunks[0] != audio_path:
                    safe_cleanup(audio_path)
            else:
                audio_chunks = [audio_path]
            log.info(f"Audio chunks: {len(audio_chunks)}")

            # ── 6. Transcribe + analyze each audio chunk ─────────
            log.info("Step 6: Transcribing + analyzing audio...")
            all_transcript = []
            for aidx, achunk in enumerate(audio_chunks):
                log.info(f"  Audio chunk {aidx+1}/{len(audio_chunks)}")
                transcript = transcribe_audio_chunk(achunk)
                if transcript:
                    all_transcript.append(transcript)
                    analysis = analyze_audio_transcript(
                        transcript, aidx, len(audio_chunks), api_key)
                    result["audio_analysis"].append({
                        "chunk_index": aidx,
                        "transcript": transcript,
                        "analysis": analysis or "ANALYSIS_FAILED",
                    })
                safe_cleanup(achunk)
            result["audio_transcript"] = "\n\n".join(all_transcript)
        else:
            result["errors"].append("Audio extraction failed")

        # ── 7. Mark success and Store in Scalable Directory ──────────
        if result["video_analysis"] or result["audio_analysis"]:
            result["status"] = "done"
        log.info(f"Video {video_id} → {result['status']}")

        # Move to scalable directory system
        channel_id = result.get("metadata", {}).get("channel_id", "unknown_channel")
        store_dir = os.path.join(os.path.dirname(work_dir), "video_store", channel_id, video_id)
        os.makedirs(store_dir, exist_ok=True)
        
        # Copy original video and yt-dlp metadata to store_dir
        if os.path.exists(dl_dir):
            for f in glob.glob(os.path.join(dl_dir, "*")):
                shutil.copy2(f, store_dir)
        
        # Save AI generated descriptions
        ai_desc_path = os.path.join(store_dir, "ai_description.json")
        with open(ai_desc_path, "w", encoding="utf-8") as f:
            json.dump({
                "video_analysis": result.get("video_analysis", []),
                "audio_analysis": result.get("audio_analysis", []),
                "audio_transcript": result.get("audio_transcript", "")
            }, f, indent=2, ensure_ascii=False)

    except Exception as e:
        log.error(f"Pipeline failed for {video_id}: {e}")
        result = None
    finally:
        safe_cleanup(vid_dir)
        # clean HF staging
        try:
            # We skip HF deletion if we actually want to keep it on HF as requested
            # hf_delete_folder(staging_prefix, hf_token)
            pass
        except Exception:
            pass

    return result

# ── results → parquet + upload ───────────────────────────────────────
def save_and_upload(results, worker_id, hf_token, work_dir):
    """Save results as parquet and upload to HF dataset."""
    import pyarrow as pa
    import pyarrow.parquet as pq

    if not results:
        log.warning("No results to save")
        return False

    rows = []
    for r in results:
        rows.append({
            "video_id": r["video_id"],
            "url": r["url"],
            "title": r.get("metadata",{}).get("title",""),
            "channel": r.get("metadata",{}).get("channel",""),
            "duration": r.get("metadata",{}).get("duration",0),
            "view_count": r.get("metadata",{}).get("view_count",0),
            "upload_date": r.get("metadata",{}).get("upload_date",""),
            "description": r.get("metadata",{}).get("description",""),
            "video_analysis": json.dumps(r.get("video_analysis",[])),
            "audio_transcript": r.get("audio_transcript",""),
            "audio_analysis": json.dumps(r.get("audio_analysis",[])),
            "hf_video_urls": json.dumps(r.get("hf_video_urls",[])),
            "status": r["status"],
            "errors": json.dumps(r.get("errors",[])),
            "worker_id": r["worker_id"],
            "processed_at": r["processed_at"],
        })

    table = pa.table({k: [row[k] for row in rows] for k in rows[0]})
    ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    fname = f"{worker_id}_{ts}.parquet"
    local_pq = os.path.join(work_dir, fname)
    pq.write_table(table, local_pq, compression="zstd")
    log.info(f"Saved {len(rows)} rows → {fname} ({os.path.getsize(local_pq)/1024:.0f} KB)")

    remote = f"{DATA_PREFIX}/{fname}"
    ok = hf_upload_file(local_pq, remote, hf_token)
    safe_cleanup(local_pq)
    if ok:
        log.info(f"✅ Uploaded to HF: {remote}")
    else:
        log.error("❌ Failed to upload parquet")
    return ok

# ── status reporting ─────────────────────────────────────────────────
def report_status(worker_id, status_data, hf_token, work_dir):
    fp = os.path.join(work_dir, "status.json")
    with open(fp, "w") as f:
        json.dump(status_data, f, indent=2)
    hf_upload_file(fp, f"status/{worker_id}.json", hf_token)
    safe_cleanup(fp)

# ── main entry ───────────────────────────────────────────────────────
def run_worker(config):
    """
    config = {
        "worker_id": "w-abc123",
        "videos": [{"video_id": "xxx", "url": "https://..."}],
        "api_keys": ["key1", "key2", ...],
        "hf_token": "hf_xxx",
        "vision_model": "Qwen/Qwen3-VL-32B-Instruct",
        "text_model": "Qwen/Qwen3-32B",
    }
    """
    worker_id = config["worker_id"]
    videos    = config["videos"]
    api_keys  = config["api_keys"]
    hf_token  = config["hf_token"]
    vision_models = config.get("vision_models")

    work_dir  = os.path.join(
        os.environ.get("KAGGLE_WORKING_DIR", "/kaggle/working"), "yt_pipeline")
    os.makedirs(work_dir, exist_ok=True)

    log.info(f"Worker {worker_id} starting — {len(videos)} videos, "
             f"{len(api_keys)} API keys, disk free: {disk_free_gb():.1f}GB")

    results = []
    status  = {
        "worker_id": worker_id,
        "total": len(videos),
        "done": 0, "failed": 0,
        "started_at": datetime.now(timezone.utc).isoformat(),
        "state": "running",
        "video_status": {},
    }

    for vidx, vinfo in enumerate(videos):
        vid_id  = vinfo["video_id"]
        vid_url = vinfo["url"]
        api_key = api_keys[vidx % len(api_keys)]  # round-robin

        log.info(f"\n[{vidx+1}/{len(videos)}] {vid_id}")

        # check disk
        if disk_free_gb() < DISK_HEADROOM_GB:
            log.error("Disk too full, stopping early")
            status["state"] = "disk_full"
            break

        result = process_video(vid_url, vid_id, api_key, hf_token,
                               worker_id, work_dir, vision_models=vision_models)

        if result is None:
            status["failed"] += 1
            status["video_status"][vid_id] = "failed"
        else:
            results.append(result)
            status["video_status"][vid_id] = result["status"]
            if result["status"] == "done":
                status["done"] += 1
            else:
                status["failed"] += 1

        # periodic save + status report
        if (vidx + 1) % 10 == 0 or vidx == len(videos) - 1:
            save_and_upload(results, worker_id, hf_token, work_dir)
            results = []  # reset batch
            status["last_update"] = datetime.now(timezone.utc).isoformat()
            report_status(worker_id, status, hf_token, work_dir)

    # final save
    if results:
        save_and_upload(results, worker_id, hf_token, work_dir)

    status["state"] = "completed"
    status["finished_at"] = datetime.now(timezone.utc).isoformat()
    report_status(worker_id, status, hf_token, work_dir)
    log.info(f"Worker {worker_id} finished: {status['done']} done, {status['failed']} failed")
    return status


def _is_notebook():
    """Detect if running inside Jupyter/Kaggle notebook."""
    try:
        return "ipykernel" in sys.modules or any("-f" in a for a in sys.argv)
    except Exception:
        return False

if __name__ == "__main__" and not _is_notebook():
    # For local testing only — skipped on Kaggle
    if len(sys.argv) > 1:
        with open(sys.argv[1]) as f:
            cfg = json.load(f)
        run_worker(cfg)
    else:
        print("Usage: python worker.py config.json")