#!/usr/bin/env python3
"""
End-to-end TRI-CMU conversion pipeline.

Source of raw episodes: S3 (populated by an external download pipeline).
Per-episode S3 layout: s3://<bucket>/cv_downloaded/TriCMU/<eid>/{video_1.mp4,
hand_pose.json, body_pose.json, pelvis_pose.csv, intrinsics.json,
annotations.json, egomotion.txt, metadata.json}

Steps:
  1. List all episode IDs present in S3 (intersect with manifest).
  2. Deterministically sample N% of them.
  3. `aws s3 cp --recursive` each sampled episode locally in parallel.
  4. Convert with anydata.converters.tri_cmu.
  5. Optionally cleanup local raw downloads.

Usage:
    aws sso login --profile sagemaker     # one-time per session
    python scripts/convert_tri_cmu.py --pct 5 --num_procs 16
"""

import argparse
import json
import os
import random
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--manifest', default='/workspace/AnyData/tri-cmu-first-batch.json')
    p.add_argument('--s3_raw',
                   default='s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_downloaded/TriCMU')
    p.add_argument('--aws_profile', default='sagemaker')
    p.add_argument('--local_raw', default='/data/cv_downloaded/TriCMU_5pct')
    p.add_argument('--local_unified', default='/data/cv_unified/TriCMU_5pct')
    p.add_argument('--pct', type=float, default=5.0,
                   help='Percentage of total dataset to sample')
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--num_procs', type=int, default=16,
                   help='Parallel workers for both download and conversion')
    p.add_argument('--episode_list', default=None,
                   help='Optional text file with episode IDs (overrides sampling)')
    p.add_argument('--episode_list_out', default='/tmp/tricmu_5pct_episodes.txt')
    p.add_argument('--skip_download', action='store_true')
    p.add_argument('--skip_convert', action='store_true')
    p.add_argument('--cleanup_raw', action='store_true',
                   help='Delete /local_raw after conversion')
    return p.parse_args()


def list_s3_episodes(s3_path, profile):
    """Return the set of episode_id directories that currently exist on S3."""
    print(f'### Listing S3 episodes under {s3_path}/ …')
    r = subprocess.run(
        ['aws', 's3', 'ls', s3_path + '/', '--profile', profile],
        capture_output=True, text=True)
    if r.returncode != 0:
        print(f'### S3 list failed:\n{r.stderr}', file=sys.stderr)
        sys.exit(1)
    eids = set()
    for line in r.stdout.splitlines():
        parts = line.strip().split()
        if len(parts) == 2 and parts[0] == 'PRE':
            eids.add(parts[1].rstrip('/'))
    print(f'### {len(eids):,} episodes currently on S3')
    return eids


def sample_episodes(manifest_path, s3_eids, pct, seed):
    with open(manifest_path) as f:
        m = json.load(f)
    all_ids = [e['episode_id'] for e in m['episodes']]
    available = sorted(set(all_ids) & s3_eids)
    n_target = max(1, int(len(all_ids) * pct / 100.0))
    n_sample = min(n_target, len(available))
    rng = random.Random(seed)
    sample = sorted(rng.sample(available, n_sample))
    print(f'### Sampled {n_sample:,} episodes ({pct}% of {len(all_ids):,} manifest, '
          f'from {len(available):,} S3-available)')
    return sample


def download_one(eid, s3_raw, local_raw, profile):
    src = f'{s3_raw}/{eid}/'
    dst = f'{local_raw}/{eid}/'
    if os.path.isdir(dst) and os.path.exists(os.path.join(dst, 'video_1.mp4')):
        return eid, True, 'cached'
    os.makedirs(dst, exist_ok=True)
    r = subprocess.run(
        ['aws', 's3', 'cp', src, dst, '--recursive',
         '--profile', profile, '--quiet'],
        capture_output=True, text=True, timeout=900)
    if r.returncode != 0:
        return eid, False, r.stderr[:200]
    return eid, True, ''


def download_all(episodes, s3_raw, local_raw, profile, num_procs):
    print(f'### Downloading {len(episodes):,} episodes → {local_raw}')
    os.makedirs(local_raw, exist_ok=True)
    done = 0
    failed = []
    with ThreadPoolExecutor(max_workers=num_procs) as pool:
        futs = [pool.submit(download_one, eid, s3_raw, local_raw, profile)
                for eid in episodes]
        for fut in as_completed(futs):
            eid, ok, msg = fut.result()
            done += 1
            if not ok:
                failed.append((eid, msg))
            if done % 100 == 0 or done == len(episodes):
                print(f'  [{done}/{len(episodes)}]  failed={len(failed)}')
    if failed:
        print(f'### {len(failed)} downloads failed. First 5:')
        for eid, msg in failed[:5]:
            print(f'    {eid}: {msg[:120]}')
    print(f'### Download done: {done - len(failed)} ok, {len(failed)} failed')


def main():
    args = parse_args()

    # 1. Determine episode list
    if args.episode_list:
        with open(args.episode_list) as f:
            episodes = [l.strip() for l in f if l.strip()]
        print(f'### Loaded {len(episodes)} episodes from {args.episode_list}')
    else:
        s3_eids = list_s3_episodes(args.s3_raw, args.aws_profile)
        episodes = sample_episodes(args.manifest, s3_eids, args.pct, args.seed)
        with open(args.episode_list_out, 'w') as f:
            for eid in episodes:
                f.write(eid + '\n')
        print(f'### Wrote {args.episode_list_out}')

    # 2. Download
    if not args.skip_download:
        download_all(episodes, args.s3_raw, args.local_raw, args.aws_profile,
                     args.num_procs)

    # 3. Convert
    if not args.skip_convert:
        # Converter computes dst = `{local_path}/{folder}/{storage}/{basename(src)}`.
        # storage = "frames" (or "videos") is appended to folder by parse_args.
        # The wrapper's `--local_unified` is interpreted as "the unified-format
        # root you want — the converter will insert /frames/ before the
        # episode-name dir". So if --local_unified=/data/cv_unified/TriCMU_5pct
        # and --local_raw=/data/cv_downloaded/TriCMU_5pct, output goes to
        # /data/cv_unified/frames/TriCMU_5pct (and is named after local_raw).
        unified_dir  = os.path.dirname(args.local_unified.rstrip('/'))   # /data/cv_unified
        unified_root = os.path.dirname(unified_dir)                      # /data
        unified_fold = os.path.basename(unified_dir)                     # cv_unified
        actual_out = (
            f'{unified_root}/{unified_fold}/frames/'
            f'{os.path.basename(args.local_raw.rstrip("/"))}'
        )
        if os.path.basename(args.local_raw.rstrip('/')) != \
           os.path.basename(args.local_unified.rstrip('/')):
            print('WARN: local_raw and local_unified basenames differ — '
                  'converter will name output after local_raw\'s basename',
                  file=sys.stderr)
        print(f'\n### Converting {args.local_raw} → {actual_out}')
        cmd = [
            sys.executable, '-m', 'anydata.converters.tri_cmu',
            args.local_raw,
            '--official',
            '--frames',
            '--num_procs', str(args.num_procs),
            '--restart',
            '--local_path', unified_root,
            '--folder', unified_fold,
        ]
        print(f'$ {" ".join(cmd)}')
        r = subprocess.run(cmd)
        if r.returncode != 0:
            print(f'### Conversion failed (exit {r.returncode})', file=sys.stderr)
            sys.exit(r.returncode)

    # 4. Cleanup
    if args.cleanup_raw:
        print(f'\n### Removing local raw: {args.local_raw}')
        subprocess.run(['rm', '-rf', args.local_raw])

    print('\n### Done.')


if __name__ == '__main__':
    main()