#!/usr/bin/env python3
"""
SNAC tokenization for FineVideo-VLA activities.

Reads final_dataset_adaptive JSONL files, extracts audio from .mp4 videos,
and tokenizes each activity segment with SNAC_24kHz in "listen" format.

Listen format — 3 tokens per SNAC base frame (base rate = 12.5 Hz → 37.5 tokens/sec):
    token_1 = codes[0][i]       + 128266        →  <snac_128266> .. <snac_132361>
    token_2 = codes[1][2*i]     + 128266 + 4096 →  <snac_132362> .. <snac_136457>
    token_3 = codes[1][2*i+1]   + 128266 +16384 →  <snac_144650> .. <snac_148745>

Same Orpheus offset scheme as MixtureVitae-Omni → tokens are directly compatible.
Total unique SNAC token strings: 3 × 4096 = 12,288.

Output: {OUTPUT_DIR}/{video_id}_snac.jsonl
    One line per activity:
    {
      "video_id":"...", "activity_id":"...", "start_sec":1.0, "end_sec":8.9,
      "has_agent": true,
      "snac_by_chunk": {
        "0": ["<snac_130055>", "<snac_133001>", "<snac_145000>", ...],  // ~9-10 tokens
        "1": [...],
        ...
      }
    }

    snac_by_chunk keys are chunk_idx (integer as string), aligned to the same 8-frame
    grid as cosmos/avclm/agent. Phase7 reads chunk_idx → snac tokens directly.

Two modes:
  --build-tasks   Scan all final_dataset_adaptive files, write snac_task_list.json.
                  Run once on login node (or task 0) before the array job.
  (default)       Load task list, process this SLURM task's slice of videos.

SLURM usage:
    SLURM_ARRAY_TASK_ID   = task index (0-based)
    SLURM_ARRAY_TASK_COUNT = total number of tasks in the array

Local test:
    python pipeline_pose/snac_finevideo.py --build-tasks
    python pipeline_pose/snac_finevideo.py  # task_id=0, num_tasks=1
"""

import argparse
import glob
import json
import logging
import math
import multiprocessing
import os
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import torch

# ── Paths ─────────────────────────────────────────────────────────────────────

VIDEO_DIR    = "/e/data1/datasets/playground/mmlaion/shared/nguyen38/videos_staging"
# Still scanned for activity time_range_sec + has_agent -- those don't depend
# on window size. chunk_timing (per-window breakdown) is IGNORED (see
# _scan_one_rank_file docstring) because this file is window=8-based (the
# pre-pivot merge) and Phase 6 hasn't rerun at window=24 yet; n_chunks is
# instead recomputed independently from CHUNK_SIZE below, same pattern as
# data_prep/omnivideo_100k/snac_omnivideo.py.
INPUT_GLOB   = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/"
                "final_dataset_adaptive/final_vla_adaptive_rank_*.jsonl")
OUTPUT_DIR   = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/snac_tokens_w24")
TASK_CACHE   = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/"
                "snac_task_list_w24.json")
HF_CACHE     = "/e/project1/reformo/nguyen38/jupiter_cache/huggingface"
SNAC_MODEL   = "hubertsiuzdak/snac_24khz"
SAMPLE_RATE  = 24000
TARGET_FPS   = 30
CHUNK_SIZE   = 24  # 2026-07-23 window=24 pivot -- must match step_a_tokenize_video.py's CHUNK_SIZE

# ── SNAC listen-format offsets (matches MixtureVitae-Omni) ───────────────────

OFFSET_L0  = 128266            # codes[0] base
OFFSET_L1A = 128266 + 4096     # codes[1] even frames   → 132362
OFFSET_L1B = 128266 + 4 * 4096 # codes[1] odd frames   → 144650

# 2026-07-23: speak-format offsets for codes[2] (fine, 50Hz, 4 sub-positions
# per base frame) -- corrected to match the REAL scheme in Huu/Chien's
# production snac_gpu.py on Leonardo (pipeline_video/snac_gpu.py), the
# Orpheus-standard SNAC packing layout. Sub-codes 0/1 sit between L1A and
# L1B (in what was wrongly assumed to be an unused gap); sub-codes 2/3 sit
# after L1B. See data_prep/laion_emotional_roleplay/tokenize_snac.py's
# OFFSET_L2 docstring for the full correction history. Must match exactly.
OFFSET_L2 = [136458, 140554, 148746, 152842]

# ── Logging ───────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
    stream=sys.stdout,
)
log = logging.getLogger(__name__)


# ─────────────────────────────────────────────────────────────────────────────
# Audio extraction
# ─────────────────────────────────────────────────────────────────────────────

def _find_ffmpeg() -> str:
    """Return path to ffmpeg binary, trying imageio_ffmpeg as fallback."""
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        return "ffmpeg"
    except (FileNotFoundError, subprocess.CalledProcessError):
        pass
    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except ImportError:
        pass
    raise RuntimeError("ffmpeg not found. Install ffmpeg or imageio_ffmpeg.")

_FFMPEG = None

def get_ffmpeg() -> str:
    global _FFMPEG
    if _FFMPEG is None:
        _FFMPEG = _find_ffmpeg()
    return _FFMPEG


def extract_full_audio(video_path: str) -> np.ndarray | None:
    """
    Extract full mono 24 kHz PCM audio from a video file.
    Pipes raw float32 PCM from ffmpeg directly to a numpy array — no temp files.
    Returns float32 array or None on any failure.
    """
    cmd = [
        get_ffmpeg(), "-y",
        "-i", video_path,
        "-vn",                     # strip video
        "-ac", "1",                # mono
        "-ar", str(SAMPLE_RATE),   # 24 kHz
        "-f",  "f32le",            # raw float32 PCM
        "-",
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=300)
        if result.returncode != 0 or not result.stdout:
            return None
        audio = np.frombuffer(result.stdout, dtype=np.float32).copy()
        return audio if len(audio) > 0 else None
    except Exception as e:
        log.debug(f"ffmpeg failed for {video_path}: {e}")
        return None


def slice_audio(
    audio: np.ndarray,
    start_sec: float,
    end_sec: float,
    sr: int = SAMPLE_RATE,
) -> np.ndarray:
    """Slice a float32 audio array to [start_sec, end_sec]."""
    s = max(0, int(start_sec * sr))
    e = min(len(audio), int(end_sec * sr))
    return audio[s:e]


# ─────────────────────────────────────────────────────────────────────────────
# SNAC tokenization
# ─────────────────────────────────────────────────────────────────────────────

def encode_listen(audio: np.ndarray, model, device: str) -> list[str]:
    """
    Encode a float32 audio array with SNAC_24kHz, return listen-format tokens.

    Listen format (3 tokens per base frame):
        <snac_{codes[0][i] + 128266}>
        <snac_{codes[1][2i]   + 132362}>
        <snac_{codes[1][2i+1] + 144650}>

    SNAC_24kHz hierarchy:
        codes[0] — base   codebook, 12.5 Hz
        codes[1] — mid    codebook, 25.0 Hz  (2× codes[0])
        codes[2] — fine   codebook, 50.0 Hz  (4× codes[0], not used in listen)

    Listen format ignores codes[2] (fine detail) to keep token count low
    (~37.5 tokens/sec vs 87.5 for full speak format). Matches MV-Omni.
    """
    tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device)  # (1,1,T)
    with torch.inference_mode():
        codes = model.encode(tensor)  # list: [codes[0], codes[1], codes[2]]

    c0 = codes[0]  # (1, N0)
    c1 = codes[1]  # (1, N1), N1 == 2*N0

    n0 = c0.shape[1]
    tokens: list[str] = []
    for i in range(n0):
        i1a = 2 * i
        i1b = 2 * i + 1
        if i1b >= c1.shape[1]:
            break  # boundary guard: shouldn't happen for valid audio
        tokens.append(f"<snac_{c0[0, i].item() + OFFSET_L0}>")
        tokens.append(f"<snac_{c1[0, i1a].item() + OFFSET_L1A}>")
        tokens.append(f"<snac_{c1[0, i1b].item() + OFFSET_L1B}>")
    return tokens


def encode_speak(audio: np.ndarray, model, device: str) -> list[str]:
    """
    Encode a float32 audio array with SNAC_24kHz, return full speak-format
    tokens (7 tokens per base frame, +133% vs encode_listen()) -- 2026-07-22
    (REPORT.md #37), decided after a real audio A/B
    (tools/snac_l2_experiment.py) showed audibly better reconstruction.

    Speak format (2026-07-23, Leo-matched order -- L2 interleaved between
    L1a and L1b, not appended after):
        <snac_{codes[0][i]     + OFFSET_L0}>
        <snac_{codes[1][2i]    + OFFSET_L1A}>
        <snac_{codes[2][4i]    + OFFSET_L2[0]}>
        <snac_{codes[2][4i+1]  + OFFSET_L2[1]}>
        <snac_{codes[1][2i+1]  + OFFSET_L1B}>
        <snac_{codes[2][4i+2]  + OFFSET_L2[2]}>
        <snac_{codes[2][4i+3]  + OFFSET_L2[3]}>
    """
    tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device)
    with torch.inference_mode():
        codes = model.encode(tensor)

    c0, c1, c2 = codes[0], codes[1], codes[2]
    n0 = c0.shape[1]
    tokens: list[str] = []
    for i in range(n0):
        i1a, i1b = 2 * i, 2 * i + 1
        i2 = [4 * i + k for k in range(4)]
        if i1b >= c1.shape[1] or i2[-1] >= c2.shape[1]:
            break
        tokens.append(f"<snac_{c0[0, i].item() + OFFSET_L0}>")
        tokens.append(f"<snac_{c1[0, i1a].item() + OFFSET_L1A}>")
        tokens.append(f"<snac_{c2[0, i2[0]].item() + OFFSET_L2[0]}>")
        tokens.append(f"<snac_{c2[0, i2[1]].item() + OFFSET_L2[1]}>")
        tokens.append(f"<snac_{c1[0, i1b].item() + OFFSET_L1B}>")
        tokens.append(f"<snac_{c2[0, i2[2]].item() + OFFSET_L2[2]}>")
        tokens.append(f"<snac_{c2[0, i2[3]].item() + OFFSET_L2[3]}>")
    return tokens


# ─────────────────────────────────────────────────────────────────────────────
# Task list building (pre-processing step)
# ─────────────────────────────────────────────────────────────────────────────

def _scan_one_rank_file(fpath: str) -> dict:
    """
    Scan one final_dataset_adaptive rank file for activity time boundaries.
    Returns {video_id: [activity_dict, ...]} for ALL activities with a valid
    time_range_sec.

    Each activity dict: activity_id, start_sec, end_sec, has_agent.

    2026-07-23: no longer reads chunk_timing (this file is the pre-window=24-
    pivot merge; its chunk_timing reflects the OLD 8-frame grid). n_chunks is
    instead recomputed in process_video() from (end_sec-start_sec) and the
    current CHUNK_SIZE=24 -- same independent-recompute pattern
    data_prep/omnivideo_100k/snac_omnivideo.py uses, so this script no longer
    needs to wait on Phase 6 merge to re-run at window=24 before it can align
    correctly.

    We tokenize ALL activities (not just agent) because:
    - Non-agent activities have seed2+cosmos → seed2+cosmos+snac trains modality transitions
    - Agent-only activities = only 14% of total; skipping the rest wastes 86% of this GPU run
    """
    tasks: dict = {}
    try:
        with open(fpath, "r", errors="replace") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    rec = json.loads(line)
                except json.JSONDecodeError:
                    continue
                vid = rec.get("video_id", "")
                if not vid:
                    continue
                for scene in rec.get("scenes", []):
                    for act in scene.get("activities", []):
                        tr = act.get("time_range_sec")
                        if not tr or len(tr) < 2:
                            continue
                        has_agent = "<agent>" in act.get("video_tokens", "")
                        tasks.setdefault(vid, []).append({
                            "activity_id": act.get("activity_id", ""),
                            "start_sec":   float(tr[0]),
                            "end_sec":     float(tr[1]),
                            "has_agent":   has_agent,
                        })
    except Exception as e:
        log.warning(f"Error scanning {fpath}: {e}")
    return tasks


def build_task_list(input_glob: str, cache_path: str, workers: int = 8) -> dict:
    """
    Scan all final_dataset_adaptive rank files in parallel to build a task list.
    Saves result to cache_path as JSON.
    Returns {video_id: [activity_dicts]}.

    This is I/O-heavy (~657 GB total) — use multiprocessing to parallelize.
    Estimated wall time: 5–15 min with 8 workers on shared filesystem.
    """
    rank_files = sorted(glob.glob(input_glob))
    if not rank_files:
        raise FileNotFoundError(f"No files matched: {input_glob}")
    log.info(f"Scanning {len(rank_files)} rank files with {workers} workers...")
    t0 = time.time()

    with multiprocessing.Pool(workers) as pool:
        partial_results = pool.map(_scan_one_rank_file, rank_files)

    # merge
    all_tasks: dict = {}
    for partial in partial_results:
        for vid, acts in partial.items():
            all_tasks.setdefault(vid, []).extend(acts)

    # deduplicate activities by activity_id (in case of overlap across ranks)
    for vid in all_tasks:
        seen = set()
        deduped = []
        for act in all_tasks[vid]:
            key = act["activity_id"]
            if key not in seen:
                seen.add(key)
                deduped.append(act)
        all_tasks[vid] = deduped

    log.info(
        f"Task list built: {len(all_tasks)} videos, "
        f"{sum(len(v) for v in all_tasks.values())} activities  "
        f"({time.time()-t0:.0f}s)"
    )
    os.makedirs(os.path.dirname(cache_path), exist_ok=True)
    with open(cache_path, "w") as f:
        json.dump(all_tasks, f)
    log.info(f"Saved task list → {cache_path}")
    return all_tasks


# ─────────────────────────────────────────────────────────────────────────────
# Chunk alignment
# ─────────────────────────────────────────────────────────────────────────────

def split_snac_by_chunks(tokens: list[str], n_chunks: int) -> dict[int, list[str]]:
    """
    Split a flat SNAC listen token list evenly across n_chunks video chunks.

    Why: SNAC rate (12.5 Hz base × 3 tokens = 37.5 tok/s) does not divide evenly
    by the video chunk rate (30fps / 8 = 3.75 Hz). Encoding the full activity once
    preserves audio context; then we split by chunk count snapping to 3-token
    boundaries (one SNAC base frame = 3 listen tokens).

    Per 8-frame chunk at 30fps (0.267s): ~3.33 SNAC base frames → 9–10 listen tokens.
    """
    n_tokens = len(tokens)
    n_base   = n_tokens // 3  # truncate to complete base frames
    tokens   = tokens[:n_base * 3]

    result: dict[int, list[str]] = {}
    for k in range(n_chunks):
        start_base = round(k       * n_base / n_chunks)
        end_base   = round((k + 1) * n_base / n_chunks)
        result[k]  = tokens[start_base * 3 : end_base * 3]
    return result


# ─────────────────────────────────────────────────────────────────────────────
# Per-video processing
# ─────────────────────────────────────────────────────────────────────────────

def process_video(
    video_id: str,
    activities: list[dict],
    model,
    device: str,
    video_dir: str,
    output_dir: str,
    skip_existing: bool,
    encode_fn=encode_listen,
) -> dict:
    """
    Tokenize all activities for one video.

    Steps:
      1. Check skip — if output file exists and skip_existing, return immediately.
      2. Extract full audio from .mp4 once (1 ffmpeg call per video).
      3. For each activity: slice audio by time_range_sec, run SNAC encode once,
         then split the flat token list across chunks (preserving audio context).
      4. Write all results to {output_dir}/{video_id}_snac.jsonl.

    Output per activity: snac_by_chunk {chunk_idx → [tokens]}
    Phase6 merge uses this directly to inject SNAC tokens per 24-frame chunk,
    aligned with the cosmos/avclm/agent tokens that fire at the same chunk
    boundaries (n_chunks recomputed from CHUNK_SIZE, not read from a file).

    Returns stats: {ok, skipped_vid, failed_audio, failed_snac, tokens}
    """
    out_path = os.path.join(output_dir, f"{video_id}_snac.jsonl")
    if skip_existing and os.path.exists(out_path):
        return {"ok": 0, "skipped_vid": len(activities), "failed_audio": 0,
                "failed_snac": 0, "tokens": 0}

    video_path = os.path.join(video_dir, f"{video_id}.mp4")
    if not os.path.exists(video_path):
        return {"ok": 0, "skipped_vid": 0, "failed_audio": len(activities),
                "failed_snac": 0, "tokens": 0}

    # Extract full audio once
    full_audio = extract_full_audio(video_path)
    if full_audio is None:
        log.warning(f"No audio: {video_path}")
        return {"ok": 0, "skipped_vid": 0, "failed_audio": len(activities),
                "failed_snac": 0, "tokens": 0}

    stats = {"ok": 0, "skipped_vid": 0, "failed_audio": 0, "failed_snac": 0, "tokens": 0}
    rows = []

    for act in activities:
        segment = slice_audio(full_audio, act["start_sec"], act["end_sec"])
        if len(segment) < int(SAMPLE_RATE * 0.1):  # skip segments < 100 ms
            stats["failed_audio"] += 1
            continue
        try:
            flat_tokens = encode_fn(segment, model, device)
        except Exception as e:
            log.warning(f"SNAC failed {video_id}/{act['activity_id']}: {e}")
            stats["failed_snac"] += 1
            continue
        if not flat_tokens:
            stats["failed_snac"] += 1
            continue

        # Split flat token list into per-chunk dicts, aligned to the same
        # 24-frame grid as cosmos/avclm/agent. n_chunks recomputed
        # independently from CHUNK_SIZE (not from a merged file's
        # chunk_timing -- see _scan_one_rank_file docstring).
        total_frames = max(1, round((act["end_sec"] - act["start_sec"]) * TARGET_FPS))
        n_chunks = math.ceil(total_frames / CHUNK_SIZE)
        by_chunk = split_snac_by_chunks(flat_tokens, n_chunks)
        snac_by_chunk = {str(k): v for k, v in by_chunk.items()}

        rows.append({
            "video_id":      video_id,
            "activity_id":   act["activity_id"],
            "start_sec":     round(act["start_sec"], 4),
            "end_sec":       round(act["end_sec"],   4),
            "has_agent":     act.get("has_agent", False),
            "snac_by_chunk": snac_by_chunk,
        })
        stats["ok"]     += 1
        stats["tokens"] += len(flat_tokens)

    if rows:
        with open(out_path, "w") as f:
            for row in rows:
                f.write(json.dumps(row) + "\n")

    return stats


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="SNAC tokenization for FineVideo-VLA")
    p.add_argument("--build-tasks",   action="store_true",
                   help="Scan final_dataset_adaptive and write snac_task_list.json, then exit.")
    p.add_argument("--input-glob",    default=INPUT_GLOB)
    p.add_argument("--output-dir",    default=OUTPUT_DIR)
    p.add_argument("--video-dir",     default=VIDEO_DIR)
    p.add_argument("--task-cache",    default=TASK_CACHE)
    p.add_argument("--hf-cache",      default=HF_CACHE)
    p.add_argument("--scan-workers",  type=int, default=8,
                   help="CPU workers for --build-tasks scan (default 8)")
    p.add_argument("--no-skip",       action="store_true",
                   help="Re-process videos even if output file exists")
    p.add_argument("--format", choices=["listen", "speak"], default="listen",
                   help="listen = L0+L1 only (current production, 3 tok/base-frame); "
                        "speak = full L0+L1+L2 (2026-07-22, 7 tok/base-frame, +133%% tokens). "
                        "speak requires L2 tokens added to the tokenizer vocab first.")
    return p.parse_args()


def main():
    args = parse_args()
    skip_existing = not args.no_skip

    # ── Set HF cache ─────────────────────────────────────────────────────────
    os.environ.setdefault("HF_HOME", args.hf_cache)
    os.makedirs(args.hf_cache,    exist_ok=True)
    os.makedirs(args.output_dir,  exist_ok=True)

    # ── Mode: build task list ─────────────────────────────────────────────────
    if args.build_tasks:
        build_task_list(args.input_glob, args.task_cache, workers=args.scan_workers)
        return

    # ── Mode: tokenize ────────────────────────────────────────────────────────

    # SLURM array vars
    task_id   = int(os.environ.get("SLURM_ARRAY_TASK_ID",    "0"))
    num_tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT",  "1"))

    # Load task list (must exist — run --build-tasks first)
    if not os.path.exists(args.task_cache):
        log.error(
            f"Task list not found: {args.task_cache}\n"
            f"Run first:  python pipeline_pose/snac_finevideo.py --build-tasks"
        )
        sys.exit(1)
    with open(args.task_cache) as f:
        all_tasks = json.load(f)

    all_vids = sorted(all_tasks.keys())
    my_vids  = all_vids[task_id::num_tasks]
    log.info(
        f"Task {task_id}/{num_tasks}: {len(my_vids)}/{len(all_vids)} videos  "
        f"skip_existing={skip_existing}"
    )

    # ── Load SNAC model ───────────────────────────────────────────────────────
    from snac import SNAC  # imported here to avoid load cost during --build-tasks
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    log.info(f"Loading SNAC model ({SNAC_MODEL}) on {device}...")
    t_load = time.time()
    os.environ["HF_HOME"] = args.hf_cache  # ensure hub/ subdir is found
    model  = SNAC.from_pretrained(SNAC_MODEL,
                                  local_files_only=True).eval().to(device)
    log.info(f"SNAC loaded ({time.time()-t_load:.1f}s)")

    # ── Process videos ────────────────────────────────────────────────────────
    cumul  = {"ok": 0, "skipped_vid": 0, "failed_audio": 0, "failed_snac": 0, "tokens": 0}
    t_start = time.time()

    encode_fn = encode_speak if args.format == "speak" else encode_listen

    for idx, vid in enumerate(my_vids, 1):
        s = process_video(
            vid, all_tasks[vid], model, device,
            args.video_dir, args.output_dir, skip_existing,
            encode_fn=encode_fn,
        )
        for k in cumul:
            cumul[k] += s[k]

        if idx % 100 == 0 or idx == len(my_vids):
            elapsed = time.time() - t_start
            rate    = idx / elapsed
            eta     = (len(my_vids) - idx) / rate if rate > 0 else 0
            log.info(
                f"[{idx:5d}/{len(my_vids)}]  vid={vid}  "
                f"ok={s['ok']} skip={s['skipped_vid']} "
                f"fail_audio={s['failed_audio']} fail_snac={s['failed_snac']}  "
                f"rate={rate:.1f}vid/s  ETA={eta/60:.0f}m  "
                f"total_tokens={cumul['tokens']:,}"
            )

    elapsed = time.time() - t_start
    log.info(
        f"DONE task {task_id}:  "
        f"ok={cumul['ok']}  skipped={cumul['skipped_vid']}  "
        f"fail_audio={cumul['failed_audio']}  fail_snac={cumul['failed_snac']}  "
        f"tokens={cumul['tokens']:,}  wall={elapsed:.0f}s"
    )


if __name__ == "__main__":
    main()