#!/usr/bin/env python3
"""
ab_pair_packager.py (list-of-pairs manifest; full model names; no file_index)
Adds pair stamps: pair_id, ab_order, ab_slot per clip.
NOW: crops both A/B to the same 30s segment per UUID using ffmpeg.

Build A/B pairs from:
  /app2/suno/data/eval_outputs/bluejay-lang-balanced-200/<uuid>/*.mp3

For each <uuid>:
  - Pick the most recent file that contains MODEL_A_SUBSTR and MODEL_B_SUBSTR.
  - Optionally shuffle which model becomes Clip A vs Clip B.
  - Crop both clips to the same 30s window (or the longest common duration if shorter).
  - Upload outputs as: <uuid>-0.mp3 (Clip A), <uuid>-1.mp3 (Clip B).
  - Write pairs.json as a list-of-pairs (Clip A / Clip B) with metadata:
      uuid, model (full), pair_id, ab_slot, ab_order.

Manifest contains ONLY url + metadata (no local path / s3_uri fields).
Local convenience files are stored as <uuid>-A.mp3 / <uuid>-B.mp3 (cropped).
"""

from __future__ import annotations
import json
import os
import random
import shutil
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from tqdm import tqdm

# =========================
# CONFIG — EDIT THESE
# =========================

EXPERIMENT_NAME = "ab6"
ROOT_DIR = Path("/app2/suno/data/eval_outputs/bluejay-lang-balanced-200")

# Use the FULL substrings that appear in filenames for each model:
MODEL_A_SUBSTR = "v3_flow_sft_t8_rd1_t28_1E6_beta100_n16_bt2_6k_last_step10"
MODEL_B_SUBSTR = "v2_infill_d4_t39_1E6_beta100_n16_bt2_acc4_3k_last"

OUTPUT_DIR = Path(f"/app2/suno/data/christian/outputs/labelmaker-{EXPERIMENT_NAME}")

# File formats to consider (lowercase, without dots)
EXTS = ["mp3"]

# Limit number of UUID pairs processed (None = no limit)
LIMIT: Optional[int] = None

# Randomize which model is A vs B per UUID
SHUFFLE_ASSIGNMENT = True
RANDOM_SEED = 0

# Number of parallel workers for processing pairs (adjust based on CPU cores)
PROCESSING_WORKERS = 8

# Prefer symlink/hardlink instead of copying; falls back to copy if links fail
# NOTE: when cropping is enabled, links are ignored (we render new cropped files).
PREFER_LINKS = True

# Remove OUTPUT_DIR before writing (use with care)
CLEAN_OUTPUT_FIRST = True

# --- Cropping ---
USE_CROP = True
CROP_DURATION_SEC = 30.0

# --- S3 upload (optional) ---
# If S3_BASE is a non-empty string, uploads will run after packaging.
S3_BASE = f"s3://suno-annotation-public/christian/labelmaker-{EXPERIMENT_NAME}"
# If you want HTTP URLs in the manifest, set the matching HTTP base:
PUBLIC_HTTP_BASE = f"https://suno-annotation-public.s3.amazonaws.com/christian/labelmaker-{EXPERIMENT_NAME}"

S3_WORKERS = 32
AWS_PROFILE: Optional[str] = None  # e.g. "default" or None

# =========================
# END CONFIG
# =========================

A_NAME = "A"
B_NAME = "B"


def find_candidates(uuid_dir: Path, model_substr: str, exts: List[str]) -> List[Path]:
    hits: List[Path] = []
    for p in uuid_dir.iterdir():
        if p.is_file() and any(p.name.lower().endswith(f".{e}") for e in exts):
            if model_substr in p.name:
                hits.append(p)
    return hits


def choose_best(cands: List[Path]) -> Optional[Path]:
    if not cands:
        return None
    # pick most recent by mtime
    return max(cands, key=lambda p: p.stat().st_mtime)


def safe_copy(src: Path, dst: Path, do_link: bool) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        dst.unlink()
    if do_link:
        try:
            os.symlink(src, dst)
            return
        except OSError:
            try:
                os.link(src, dst)
                return
            except OSError:
                pass
    shutil.copy2(src, dst)


def ffprobe_duration_seconds(path: Path) -> Optional[float]:
    """Return duration in seconds using ffprobe, or None on failure."""
    try:
        cmd = [
            "ffprobe",
            "-v",
            "quiet",
            "-show_entries",
            "format=duration",
            "-of",
            "csv=p=0",
            str(path),
        ]
        res = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return float(res.stdout.strip())
    except Exception:
        return None


def ffmpeg_crop(
    input_path: Path,
    output_path: Path,
    start_sec: float,
    dur_sec: float,
    loudness_normalize: bool = True,
    target_lufs: float = -14.0,
) -> bool:
    """
    Crop audio to [start_sec, start_sec + dur_sec] using ffmpeg, then encode to MP3 (192k).
    - Cuts to PCM WAV first for sample-accurate edits (avoids MP3 encoder delay artifacts).
    - Adds 5 ms fades at in/out to prevent clicks on hard cuts.
    - If loudness_normalize, applies EBU R128 loudnorm with a bit more headroom (TP=-3.0).
    Returns True on success, False on failure.
    """
    try:
        if dur_sec <= 0:
            return False

        output_path.parent.mkdir(parents=True, exist_ok=True)
        if output_path.exists():
            output_path.unlink()

        # Temporary WAV path in the same directory (avoids cross-device temp issues)
        tmp_wav = output_path.with_suffix(".fftmp.wav")

        # Build audio filter chain
        fades = "afade=t=in:ss=0:d=0.005,afade=t=out:st=end-0.005:d=0.005"
        if loudness_normalize:
            # More headroom than your original to prevent codec-induced overs
            # You can tweak LRA if you want less movement on short clips.
            af = f"loudnorm=I={target_lufs}:TP=-3.0:LRA=11,{fades}"
        else:
            af = fades

        end_sec = start_sec + dur_sec

        # 1) Accurate crop to PCM WAV (sample-accurate since we re-encode to PCM)
        cmd_wav = [
            "ffmpeg",
            "-y",
            "-i",
            str(input_path),
            "-ss",
            f"{start_sec:.3f}",
            "-to",
            f"{end_sec:.3f}",
            "-af",
            af,
            "-c:a",
            "pcm_s16le",
            str(tmp_wav),
        ]

        # Silence ffmpeg during normal runs; comment these to see warnings for debugging
        subprocess.run(
            cmd_wav, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
        )

        # 2) Encode cleanly to MP3 (separate step avoids cut artifacts from MP3 encoder delay)
        cmd_mp3 = [
            "ffmpeg",
            "-y",
            "-i",
            str(tmp_wav),
            "-c:a",
            "libmp3lame",
            "-b:a",
            "192k",
            str(output_path),
        ]
        subprocess.run(
            cmd_mp3, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
        )

        # Cleanup
        try:
            os.remove(tmp_wav)
        except OSError:
            pass

        return True

    except Exception:
        # Best-effort cleanup if something failed mid-way
        try:
            if os.path.exists(tmp_wav):
                os.remove(tmp_wav)
        except Exception:
            pass
        return False


def make_numeric_name(uuid: str, index: int, src: Path) -> str:
    # e.g., "<uuid>-0.mp3" or "<uuid>-1.mp3"
    return f"{uuid}-{index}{src.suffix.lower()}"


def collect_pairs(
    root: Path, model_a: str, model_b: str, exts: List[str]
) -> List[Tuple[str, Path, Path]]:
    pairs: List[Tuple[str, Path, Path]] = []
    for uuid_dir in sorted(root.iterdir()):
        if not uuid_dir.is_dir():
            continue
        uuid = uuid_dir.name
        a_cands = find_candidates(uuid_dir, model_a, exts)
        b_cands = find_candidates(uuid_dir, model_b, exts)
        a = choose_best(a_cands)
        b = choose_best(b_cands)
        if a and b:
            pairs.append((uuid, a, b))
    return pairs


def aws_s3_cp(
    local: Path, s3_uri: str, aws_profile: Optional[str] = None
) -> Tuple[str, bool, str]:
    """Upload a file to S3 using AWS CLI. Returns (s3_uri, ok, msg)."""
    env = os.environ.copy()
    if aws_profile:
        env["AWS_PROFILE"] = aws_profile

    # Resolve symlinks to avoid CLI edge-cases
    try:
        if local.is_symlink():
            local = local.resolve()
    except Exception:
        pass

    cmd = ["aws", "s3", "cp", str(local), s3_uri]
    try:
        subprocess.run(
            cmd,
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            env=env,
        )
        return (s3_uri, True, "")
    except subprocess.CalledProcessError as e:
        return (s3_uri, False, f"aws s3 cp failed: {e}")
    except FileNotFoundError:
        return (s3_uri, False, "aws CLI not found on PATH")


def process_single_pair(args) -> Tuple[bool, List[Dict], List[Tuple[Path, str]]]:
    """Process a single pair and return success status, manifest data, and upload jobs."""
    uuid, src_a, src_b, left_src, right_src, left_model, right_model = args

    # Determine aligned crop window
    if USE_CROP:
        dur_left = ffprobe_duration_seconds(left_src)
        dur_right = ffprobe_duration_seconds(right_src)

        if dur_left is None or dur_right is None:
            start_sec = 0.0
            dur_sec = min(
                CROP_DURATION_SEC,
                (dur_left or CROP_DURATION_SEC),
                (dur_right or CROP_DURATION_SEC),
            )
            do_crop = False
        else:
            common = min(dur_left, dur_right)
            if common <= 0.05:
                start_sec = 0.0
                dur_sec = min(CROP_DURATION_SEC, common)
                do_crop = False
            else:
                if common <= CROP_DURATION_SEC:
                    start_sec = 0.0
                    dur_sec = common
                else:
                    max_start = max(0.0, common - CROP_DURATION_SEC)
                    start_sec = random.uniform(0.0, max_start)
                    dur_sec = CROP_DURATION_SEC
                do_crop = True
    else:
        start_sec = 0.0
        dur_sec = CROP_DURATION_SEC
        do_crop = False

    # Local output files
    uuid_out = OUTPUT_DIR / uuid
    a_dst_local = uuid_out / f"{uuid}-A{left_src.suffix.lower()}"
    b_dst_local = uuid_out / f"{uuid}-B{right_src.suffix.lower()}"

    if USE_CROP and do_crop:
        ok_a = ffmpeg_crop(left_src, a_dst_local, start_sec, dur_sec)
        ok_b = ffmpeg_crop(right_src, b_dst_local, start_sec, dur_sec)
        if not (ok_a and ok_b):
            safe_copy(left_src, a_dst_local, PREFER_LINKS)
            safe_copy(right_src, b_dst_local, PREFER_LINKS)
    else:
        safe_copy(left_src, a_dst_local, PREFER_LINKS)
        safe_copy(right_src, b_dst_local, PREFER_LINKS)

    # S3 keys
    a_key = make_numeric_name(uuid, 0, left_src)
    b_key = make_numeric_name(uuid, 1, right_src)

    # URLs for manifest
    if PUBLIC_HTTP_BASE:
        http_base = PUBLIC_HTTP_BASE.rstrip("/")
        a_url = f"{http_base}/{a_key}"
        b_url = f"{http_base}/{b_key}"
    else:
        a_url = ""
        b_url = ""

    # Manifest data
    pair_id = uuid
    ab_order_str = f"A={left_model};B={right_model}"

    clip_a = {
        "key": f"{uuid}_0",
        "name": "Clip A",
        "url": a_url,
        "metadata": {
            "uuid": uuid,
            "model": left_model,
            "pair_id": pair_id,
            "ab_slot": "A",
            "ab_order": ab_order_str,
            "start_sec": start_sec,
            "dur_sec": dur_sec,
        },
    }
    clip_b = {
        "key": f"{uuid}_1",
        "name": "Clip B",
        "url": b_url,
        "metadata": {
            "uuid": uuid,
            "model": right_model,
            "pair_id": pair_id,
            "ab_slot": "B",
            "ab_order": ab_order_str,
        },
    }

    upload_jobs = []
    if S3_BASE:
        base = S3_BASE.rstrip("/")
        upload_jobs = [
            (a_dst_local, f"{base}/{a_key}"),
            (b_dst_local, f"{base}/{b_key}"),
        ]

    return True, [clip_a, clip_b], upload_jobs


def do_parallel_uploads(
    jobs: List[Tuple[Path, str]], workers: int, aws_profile: Optional[str]
) -> Dict[str, bool]:
    results: Dict[str, bool] = {}
    with ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {
            ex.submit(aws_s3_cp, local, s3, aws_profile): (local, s3)
            for (local, s3) in jobs
        }
        for fut in as_completed(futs):
            _, s3 = futs[fut]
            try:
                s3_uri, ok, _ = fut.result()
            except Exception:
                ok = False
                s3_uri = s3
            results[s3_uri] = ok
    return results


def main() -> None:
    # Validate inputs
    if not ROOT_DIR.exists():
        print(f"[ERR] Root does not exist: {ROOT_DIR}", file=sys.stderr)
        sys.exit(1)

    if CLEAN_OUTPUT_FIRST and OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    exts = [e.lower().lstrip(".") for e in EXTS]

    pairs = collect_pairs(ROOT_DIR, MODEL_A_SUBSTR, MODEL_B_SUBSTR, exts)
    if LIMIT is not None:
        pairs = pairs[:LIMIT]

    if not pairs:
        print(
            "[WARN] No pairs found. Check model substrings and EXTS.", file=sys.stderr
        )
        sys.exit(2)

    print(f"[INFO] Found {len(pairs)} candidate pairs")

    if SHUFFLE_ASSIGNMENT:
        random.seed(RANDOM_SEED)

    # Prepare arguments for parallel processing
    process_args = []
    for uuid, src_a, src_b in pairs:
        # Decide A/B assignment
        if SHUFFLE_ASSIGNMENT and random.random() < 0.5:
            left_src, right_src = src_b, src_a
            left_model, right_model = MODEL_B_SUBSTR, MODEL_A_SUBSTR
        else:
            left_src, right_src = src_a, src_b
            left_model, right_model = MODEL_A_SUBSTR, MODEL_B_SUBSTR

        process_args.append(
            (uuid, src_a, src_b, left_src, right_src, left_model, right_model)
        )

    # Process pairs in parallel
    manifest: List[List[Dict]] = []
    upload_jobs: List[Tuple[Path, str]] = []
    kept = 0

    # Use configured number of workers for CPU-intensive tasks (ffmpeg/ffprobe)
    max_workers = min(PROCESSING_WORKERS, len(process_args))

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks
        future_to_args = {
            executor.submit(process_single_pair, args): args for args in process_args
        }

        # Process completed tasks
        for future in tqdm(
            as_completed(future_to_args),
            total=len(process_args),
            desc="Processing pairs",
        ):
            try:
                success, pair_manifest, pair_uploads = future.result()
                if success:
                    manifest.append(pair_manifest)
                    upload_jobs.extend(pair_uploads)
                    kept += 1
            except Exception as e:
                print(f"[WARN] Failed to process pair: {e}", file=sys.stderr)

    # Write manifest
    manifest_path = OUTPUT_DIR / "pairs.json"
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    print(f"[DONE] Prepared {kept} pairs in {OUTPUT_DIR}")
    print(f"[DONE] Manifest: {manifest_path}")

    # Perform S3 uploads if requested
    if S3_BASE and upload_jobs:
        print(
            f"[INFO] Uploading {len(upload_jobs)} files to S3 with {S3_WORKERS} workers..."
        )
        results = do_parallel_uploads(upload_jobs, S3_WORKERS, AWS_PROFILE)
        failed = [uri for uri, ok in results.items() if not ok]
        if failed:
            print(
                f"[WARN] {len(failed)} uploads failed. Example: {failed[:3]}",
                file=sys.stderr,
            )
            sys.exit(3)
        print("[DONE] S3 upload complete.")


if __name__ == "__main__":
    main()