"""
Download TRI-CMU dataset in batches, upload to S3, delete locally.

Downloads episodes from dl.mecka.ai URLs in tri-cmu-first-batch.json,
stores per-episode directories matching EgoVerse's cv_downloaded layout,
uploads each batch to S3, then deletes the local copy.

Per-episode structure:
  <episode_id>/
    video_1.mp4
    egomotion.txt
    hand_pose.json
    body_pose.json
    pelvis_pose.csv
    intrinsics.json
    annotations.json
    metadata.json          ← episode metadata from the manifest

Usage:
    # Download + upload batch 0 (first 5%)
    python scripts/download_tri_cmu.py --batch 0 --upload

    # Download batch 3 with 64 workers, no upload
    python scripts/download_tri_cmu.py --batch 3 --num_workers 64

    # Download all batches sequentially
    python scripts/download_tri_cmu.py --batch all --upload
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock, Event

# File key → extension mapping
FILE_KEYS = {
    'video_1': 'mp4',
    'egomotion': 'txt',
    'hand_pose': 'json',
    'body_pose': 'json',
    'pelvis_pose': 'csv',
    'intrinsics': 'json',
    'annotations': 'json',
}

S3_BUCKET = 's3://tri-ml-sandbox-16011-us-west-2-datasets/cv_downloaded/TriCMU'

print_lock = Lock()
sso_event = Event()
sso_event.set()  # starts as "SSO is valid"


def safe_print(msg):
    with print_lock:
        print(msg, flush=True)


def check_sso_valid():
    """Quick check if SSO is still valid."""
    result = subprocess.run(
        ['aws', 'sts', 'get-caller-identity', '--profile', 'sagemaker'],
        capture_output=True, timeout=10,
    )
    return result.returncode == 0


def wait_for_sso():
    """Block until SSO is refreshed. Only one thread triggers the prompt."""
    if sso_event.is_set():
        # First thread to notice: clear the event and prompt user
        sso_event.clear()
        safe_print('\n' + '=' * 60)
        safe_print('### SSO SESSION EXPIRED')
        safe_print('### Please run:  aws sso login --profile sagemaker')
        safe_print('### Uploads are paused. Downloads continue.')
        safe_print('=' * 60 + '\n')
        # Poll until SSO is back
        while not check_sso_valid():
            time.sleep(30)
        safe_print('### SSO refreshed — resuming uploads')
        sso_event.set()
    else:
        # Another thread already triggered; just wait
        sso_event.wait()


def download_file(url, dst_path, retries=3):
    """Download a single file with curl. Returns True on success."""
    for attempt in range(retries):
        try:
            result = subprocess.run(
                ['curl', '-s', '-f', '-o', dst_path, url],
                capture_output=True, timeout=600,
            )
            if result.returncode == 0 and os.path.exists(dst_path):
                return True
        except subprocess.TimeoutExpired:
            pass
        # Clean up partial file
        if os.path.exists(dst_path):
            os.remove(dst_path)
        if attempt < retries - 1:
            time.sleep(1)
    return False


def download_episode(episode, local_base):
    """Download all files for one episode. Returns (episode_id, success, error_msg)."""
    ep_id = episode['episode_id']
    ep_dir = os.path.join(local_base, ep_id)
    os.makedirs(ep_dir, exist_ok=True)

    # Download each file
    for key, ext in FILE_KEYS.items():
        url = episode.get(key)
        if not url:
            continue
        dst = os.path.join(ep_dir, f'{key}.{ext}')
        if os.path.exists(dst):
            continue  # resume support
        if not download_file(url, dst):
            # Clean up partial episode
            shutil.rmtree(ep_dir, ignore_errors=True)
            return ep_id, False, f'failed to download {key}'

    # Save episode metadata (minus download URLs)
    meta = {k: v for k, v in episode.items() if k not in FILE_KEYS}
    with open(os.path.join(ep_dir, 'metadata.json'), 'w') as f:
        json.dump(meta, f, indent=2)

    return ep_id, True, None


def upload_and_delete_episode(ep_id, local_base, s3_base):
    """Upload one episode to S3 and delete locally. Returns (ep_id, success).

    On SSO expiry, pauses and waits for re-auth before retrying.
    """
    ep_dir = os.path.join(local_base, ep_id)
    s3_dir = f'{s3_base}/{ep_id}/'

    for attempt in range(3):
        sso_event.wait()
        result = subprocess.run(
            ['aws', 's3', 'cp', ep_dir, s3_dir,
             '--recursive', '--profile', 'sagemaker', '--quiet'],
            capture_output=True, timeout=600,
        )
        if result.returncode == 0:
            shutil.rmtree(ep_dir, ignore_errors=True)
            return ep_id, True

        output = (result.stderr.decode() + result.stdout.decode()).lower()
        if 'expired' in output or 'invalid' in output or 'sso' in output:
            wait_for_sso()
            continue
        safe_print(f'### UPLOAD FAILED {ep_id}: {output[:200]}')
        return ep_id, False

    safe_print(f'### UPLOAD FAILED {ep_id} after 3 attempts')
    return ep_id, False


def process_batch(episodes, args):
    """Download, upload, and delete a batch of episodes."""
    local_base = args.local_path
    os.makedirs(local_base, exist_ok=True)

    total = len(episodes)
    downloaded = 0
    uploaded = 0
    failed_dl = 0
    failed_ul = 0
    t0 = time.time()

    # Pipeline: download and upload in parallel with separate thread pools.
    dl_pool = ThreadPoolExecutor(max_workers=args.num_workers)
    ul_pool = ThreadPoolExecutor(max_workers=args.num_upload_workers)

    upload_futures = []

    def submit_upload(ep_id):
        if args.upload:
            fut = ul_pool.submit(upload_and_delete_episode, ep_id, local_base, S3_BUCKET)
            upload_futures.append(fut)

    # Submit all downloads
    dl_futures = {
        dl_pool.submit(download_episode, ep, local_base): ep['episode_id']
        for ep in episodes
    }

    for future in as_completed(dl_futures):
        ep_id, success, err = future.result()
        if success:
            downloaded += 1
            submit_upload(ep_id)
            if downloaded % 500 == 0:
                elapsed = time.time() - t0
                rate = downloaded / elapsed
                eta_h = (total - downloaded) / rate / 3600 if rate > 0 else 0
                safe_print(f'### Downloaded {downloaded}/{total} '
                           f'({elapsed/60:.0f}min, {rate:.1f} eps/s, '
                           f'ETA {eta_h:.1f}h), dl_failed={failed_dl}')
        else:
            failed_dl += 1
            safe_print(f'### DOWNLOAD FAILED {ep_id}: {err}')

    dl_pool.shutdown(wait=True)
    safe_print(f'### All downloads done: {downloaded}/{total} '
               f'in {(time.time()-t0)/60:.0f}min')

    # Wait for remaining uploads
    for fut in as_completed(upload_futures):
        ep_id, success = fut.result()
        if success:
            uploaded += 1
        else:
            failed_ul += 1

    ul_pool.shutdown(wait=True)

    elapsed = time.time() - t0
    safe_print(f'### BATCH DONE in {elapsed/60:.0f}min: '
               f'downloaded={downloaded}, uploaded={uploaded}, '
               f'dl_failed={failed_dl}, ul_failed={failed_ul}')

    # Report any episodes left locally (upload failures)
    remaining = os.listdir(local_base) if os.path.isdir(local_base) else []
    if remaining:
        safe_print(f'### {len(remaining)} episodes remain locally (upload failures)')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--manifest', type=str,
                        default='/workspace/AnyData/tri-cmu-first-batch.json')
    parser.add_argument('--batch', type=str, required=True,
                        help='Batch index (0-19) or "all"')
    parser.add_argument('--batch_pct', type=float, default=5.0,
                        help='Batch size as percentage of total (default: 5%%)')
    parser.add_argument('--num_workers', type=int, default=64,
                        help='Parallel download workers')
    parser.add_argument('--num_upload_workers', type=int, default=64,
                        help='Parallel upload workers')
    parser.add_argument('--local_path', type=str,
                        default='/data/cv_downloaded/TriCMU')
    parser.add_argument('--start_batch', type=str, default=None,
                        help='Starting batch index when using --batch all (skip earlier batches)')
    parser.add_argument('--start_idx', type=int, default=None,
                        help='Override the start manifest index of the first batch '
                             '(useful to skip partway through a batch when resuming)')
    parser.add_argument('--upload', action='store_true',
                        help='Upload to S3 and delete locally after each episode')
    args = parser.parse_args()

    with open(args.manifest) as f:
        manifest = json.load(f)

    episodes = manifest['episodes']
    total = len(episodes)
    batch_size = int(total * args.batch_pct / 100)

    safe_print(f'### Total episodes: {total}')
    safe_print(f'### Batch size ({args.batch_pct}%): {batch_size}')
    safe_print(f'### Workers: {args.num_workers} download, '
               f'{args.num_upload_workers} upload')

    if args.batch == 'all':
        start = int(args.start_batch) if args.start_batch else 0
        batches = range(start, int(100 / args.batch_pct))
    else:
        batches = [int(args.batch)]

    first_batch = True
    for batch_idx in batches:
        start = batch_idx * batch_size
        end = min(start + batch_size, total)
        if first_batch and args.start_idx is not None:
            start = max(start, args.start_idx)
        first_batch = False
        if start >= total or start >= end:
            continue
        batch_episodes = episodes[start:end]
        safe_print(f'### BATCH {batch_idx}: episodes {start}-{end} '
                    f'({len(batch_episodes)} episodes)')
        process_batch(batch_episodes, args)


if __name__ == '__main__':
    main()