# Copyright 2026 Toyota Research Institute.  All rights reserved.
# TODO(bvh): consider refactoring to use boto3 instead of subprocess aws CLI calls?
# TODO(bvh): unify with custom/utils/aws.py in cpred2 (shared retry/robust logic).
# NOTE(bvh): keep aws_s3_cp/sync/rm behavior consistent with custom/utils/aws.py.

import os
import time
import shutil
import subprocess
import tarfile
import traceback
from typing import List
import numpy as np
import multiprocessing as mp
import boto3
from glob import glob
from rich.console import Console

_console = Console(stderr=True)
_RETRY_COUNT = 10
_RETRY_DELAY = 5  # seconds


def _err_detail(result: subprocess.CompletedProcess) -> str:
    '''
    Combine stdout + stderr into a single non-empty error string.
    aws s3 cp/sync/rm route per-file errors to STDOUT (not stderr); ignoring stdout
    swallowed the real AccessDenied/etc reason. Always prefer stderr when present
    but fall back to stdout so the underlying cli message is preserved.
    '''
    err = (result.stderr or '').strip()
    out = (result.stdout or '').strip()
    if err and out:
        return f'{err}\n{out}'
    return err or out or '(no aws cli output captured)'


def _warn_and_traceback(msg: str) -> None:
    # Format the entire warning + stack as ONE string and emit via a single
    # _console.print call so that tqdm (also on stderr, using \r without \n
    # in non-tty mode) cannot interleave between the frames.
    stack = ''.join(traceback.format_stack())
    block = (
        f'\n{"=" * 60}\n'
        f'[AWS WARNING] {msg}\n'
        f'Stack trace (caller context):\n'
        f'{stack}'
        f'{"=" * 60}\n'
    )
    _console.print(block, style='bold red', markup=False, highlight=False)


####################################
#  Shell-free helpers for common I/O
####################################


def aws_s3_cp(src, dst, quiet=True, extras='', check=True, robust=False):
    """Copy a single file to/from S3.

    Always retries up to _RETRY_COUNT times on failure.
    robust=False (default): raise RuntimeError on final failure (unless check=False).
    robust=True: on final failure, print a loud warning + stack trace but return False.
    """
    cmd = f'aws s3 cp {src} {dst}'
    if quiet:
        cmd += ' --only-show-errors'  # was --quiet; --quiet suppresses cli error lines too
    if extras:
        cmd += f' {extras}'
    result = None
    for attempt in range(1, _RETRY_COUNT + 1):
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        if result.returncode == 0:
            return True
        if attempt < _RETRY_COUNT:
            _console.print(
                f'[yellow][AWS] aws s3 cp attempt {attempt}/{_RETRY_COUNT} failed, '
                f'retrying in {_RETRY_DELAY}s...\n  {src} -> {dst}\n  '
                f'{_err_detail(result)}[/yellow]'
            )
            time.sleep(_RETRY_DELAY)
    err_msg = f'aws s3 cp failed ({_RETRY_COUNT} attempt(s)): {src} -> {dst}\n{_err_detail(result)}'
    if robust:
        _warn_and_traceback(err_msg)
        return False
    if check:
        raise RuntimeError(err_msg)
    return False


def aws_s3_sync(src, dst, quiet=True, extras='', robust=False):
    """Sync directories to/from S3.

    Always retries up to _RETRY_COUNT times on failure.
    robust=False (default): raise RuntimeError on final failure.
    robust=True: on final failure, print a loud warning + stack trace but return False.
    """
    cmd = f'aws s3 sync {src} {dst}'
    if quiet:
        cmd += ' --only-show-errors'  # was --quiet; --quiet suppresses cli error lines too
    if extras:
        cmd += f' {extras}'
    result = None
    for attempt in range(1, _RETRY_COUNT + 1):
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        if (not quiet or 'dryrun' in extras) and result.stdout != '':
            print(result.stdout)
        if result.returncode == 0:
            return True
        if attempt < _RETRY_COUNT:
            _console.print(
                f'[yellow][AWS] aws s3 sync attempt {attempt}/{_RETRY_COUNT} failed, '
                f'retrying in {_RETRY_DELAY}s...\n  {src} -> {dst}\n  '
                f'{_err_detail(result)}[/yellow]'
            )
            time.sleep(_RETRY_DELAY)
    err_msg = f'aws s3 sync failed ({_RETRY_COUNT} attempt(s)): {src} -> {dst}\n{_err_detail(result)}'
    if robust:
        _warn_and_traceback(err_msg)
        return False
    raise RuntimeError(err_msg)


def aws_s3_exists(s3_path):
    """Return True if a single S3 object exists (via aws s3 ls), False otherwise."""
    result = subprocess.run(
        f'aws s3 ls {s3_path}', shell=True, capture_output=True,
    )
    return result.returncode == 0 and len(result.stdout.strip()) > 0


def aws_s3_rm(path, recursive=False, quiet=True, extras='', robust=False):
    """Remove objects from S3.

    Always retries up to _RETRY_COUNT times on failure.
    robust=False (default): raise RuntimeError on final failure.
    robust=True: on final failure, print a loud warning + stack trace but return False.
    """
    cmd = f'aws s3 rm {path}'
    if recursive:
        cmd += ' --recursive'
    if quiet:
        cmd += ' --only-show-errors'  # was --quiet; --quiet suppresses cli error lines too
    if extras:
        cmd += f' {extras}'
    result = None
    for attempt in range(1, _RETRY_COUNT + 1):
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        if result.returncode == 0:
            return True
        if attempt < _RETRY_COUNT:
            _console.print(
                f'[yellow][AWS] aws s3 rm attempt {attempt}/{_RETRY_COUNT} failed, '
                f'retrying in {_RETRY_DELAY}s...\n  {path}\n  '
                f'{_err_detail(result)}[/yellow]'
            )
            time.sleep(_RETRY_DELAY)
    err_msg = f'aws s3 rm failed ({_RETRY_COUNT} attempt(s)): {path}\n{_err_detail(result)}'
    if robust:
        _warn_and_traceback(err_msg)
        return False
    raise RuntimeError(err_msg)


def rm_sync(path, args):
    if isinstance(path, list):
        for p in path:
            rm_sync(p, args)
    else:
        if os.path.exists(path):
            remove_path(path)
        if args.upload:
            path_s3 = path.replace(args.local_path, args.s3_path)
            aws_s3_rm(path_s3, robust=True)


def mv_sync(path, dst_dir, args):
    '''
    Move a file into dst_dir locally and (if uploading) on S3, instead of deleting it.
    '''
    if isinstance(path, list):
        for p in path:
            mv_sync(p, dst_dir, args)
        return
    dst = f'{dst_dir}/{os.path.basename(path)}'
    if os.path.exists(path):
        os.makedirs(dst_dir, exist_ok=True)
        if os.path.exists(dst):  # idempotent on reruns: shutil.move into an existing dst would nest/raise
            remove_path(dst)
        shutil.move(path, dst)
    if args.upload:
        path_s3 = path.replace(args.local_path, args.s3_path)
        dst_s3 = dst.replace(args.local_path, args.s3_path)
        # upload the moved local file directly (robust even if the source was never on S3);
        # fall back to an S3-to-S3 copy when the file only exists remotely
        moved = aws_s3_cp(dst, dst_s3, robust=True) if os.path.exists(dst) \
            else aws_s3_cp(path_s3, dst_s3, robust=True)
        if moved:
            aws_s3_rm(path_s3, robust=True)


def create_tar(archive_path, src_dir):
    """Create a tar.gz archive from the contents of *src_dir*."""
    os.makedirs(os.path.dirname(archive_path), exist_ok=True)
    with tarfile.open(archive_path, 'w:gz') as tar:
        for name in sorted(os.listdir(src_dir)):
            tar.add(os.path.join(src_dir, name), arcname=name)


def remove_path(path):
    """Remove a file or directory tree (no error if missing)."""
    if os.path.isdir(path):
        shutil.rmtree(path)
    elif os.path.exists(path):
        os.remove(path)


def extract_tar(archive_path, dst_dir, delete=True):
    """Extract a tar.gz archive into *dst_dir* and optionally delete the archive."""
    os.makedirs(dst_dir, exist_ok=True)
    with tarfile.open(archive_path, 'r:gz') as tar:
        tar.extractall(dst_dir)
    if delete:
        os.remove(archive_path)


def strip_s3_prefix(paths, s3_url):
    """Convert bucket-relative paths (from list_s3_recursive) to paths relative to *s3_url*."""
    parts = s3_url.replace('s3://', '').split('/', 1)
    prefix = (parts[1] + '/') if len(parts) > 1 and parts[1] else ''
    return [p[len(prefix):] for p in paths if not prefix or p.startswith(prefix)]

####################################

def get_splits(seqs, args):
    n_seqs = len(seqs)
    if n_seqs < args.num_procs:
        args.num_procs = len(seqs)
    splits = np.linspace(0, n_seqs, args.num_procs + 1)
    return [int(s) for s in splits]


def get_local_files_and_folders(local_path):
    data = sorted(glob(local_path + '/*'))
    files = [f for f in data if os.path.isfile(f)]
    folders = [f for f in data if os.path.isdir(f)]
    return files, folders    


def get_s3_files_and_folders(s3_path):
    proc = subprocess.Popen(f'aws s3 ls {s3_path}/', stdout=subprocess.PIPE, shell=True)
    data = proc.communicate()[0].decode()
    data = data.replace('\n', ' ').split(' ')
    data = [o for o in data if len(o) > 0]
    data = [o for o in data if o != 'PRE']
    proc.kill()

    extensions = ['txt', 'md', 'json', 'pkl', 'npy', 'npz', 'yaml', 'zip', 'tar', 'gz']

    files, folders = [], []
    for i in range(len(data)):
        if data[i].endswith('/'):
            folders.append(data[i][:-1])
        else:
            for ext in extensions:
                if data[i].endswith(f'.{ext}'):
                    files.append(data[i])
                    break

    folders = [f'{s3_path}/{folder}' for folder in folders]
    files = [f'{s3_path}/{file}' for file in files]

    return files, folders


def get_extras_base(args):
    extras = ''
    if not args.verbose and not args.dryrun:
        extras += ' --quiet'
    if args.dryrun:
        extras += ' --dryrun'
    return extras


def get_seqs_subset(seqs, subset):
    """Filter sequences to get just a subset"""
    if subset is None: return seqs # No subset, return all
    seqs = sorted(seqs) # Sort seqs to ensure the same ones are always picked
    if isinstance(subset, str) and '/' in subset: 
        ### Example: 10/2 breaks the sequences into 10 [0,1,2,3,...) parts and process only part 2
        if '//' in subset:  # Apply subset of subsets (for faster parallel processing)
            subset1, subset2 = subset.split('//')
            skip, start = subset1.split('/')
            start, skip = int(start), int(skip)
            seqs = seqs[start::skip]
            skip, start = subset2.split('/')
            start, skip = int(start), int(skip)
            seqs = seqs[start::skip]
        else:
            skip, start = subset.split('/')
            start, skip = int(start), int(skip)
            seqs = seqs[start::skip]
    else:
        ### Get a fixed number of sequences (linearly spaced)
        seqs = [seqs[int(i)] for i in np.linspace(0, len(seqs) - 1, int(subset))]
    return seqs

####################################

def multi_thread(seqs, files, args, name, fn_seqs, fn_files=None, tarfiles=False):

    # There are no sequences, so just return
    if len(seqs) == 0:
        return 

    # Source string to show
    if hasattr(args, 'src'):
        if hasattr(args, 'task') and args.task is not None:
            src_show = args.task
        else:
            src_show = os.path.dirname(seqs[0]) if not tarfiles else args.src

    # Destination string to show
    if hasattr(args, 'dst'):
        dst_show = args.dst

    # Subset being used
    subset = args.subset if hasattr(args, 'subset') else None

    # Screen prints when starting
    print(f'#################################################')
    print(f'###### {name}')
    print(f'#################################################')
    if seqs is not None:
        print(f'### NUM SEQUENCES: {len(seqs)}')
    if files is not None:
        print(f'### NUM FILES: {len(files)}')
    if subset is not None:
        print(f'### SUBSET: {args.subset}')
    print(f'#################################################')
    if hasattr(args, 'src'):
        print(f'### SRC: {src_show}')
    if hasattr(args, 'dst'):
        print(f'### DST: {dst_show}')
    if hasattr(args, 'extras'):
        print(f'### ARGS: {args.extras}')
    print(f'#################################################')

    # Multi-thread processing for sequences
    if args.num_procs > 0:
        jobs: List[mp.Process] = []
        splits = get_splits(seqs, args)
        for i in range(args.num_procs):
            jobs.append(mp.Process(target=fn_seqs, args=(i, seqs[splits[i]:splits[i+1]], args)))
        for j in jobs:
            j.start()
        for j in jobs:
            j.join()
    else:
        fn_seqs(0, seqs, args)

    # Download files now
    if fn_files is not None:
        fn_files(files, args.dst)

    # Screen prints when finishing
    if not hasattr(args, 'quiet') or (hasattr(args, 'quiet') and not args.quiet):
        print(f'#################################################')
        print(f'###### {name} DONE!!!')
        print(f'#################################################')
        if seqs is not None:
            print(f'### NUM SEQUENCES: {len(seqs)}')
        if files is not None:
            print(f'### NUM FILES: {len(files)}')
        if subset is not None:
            print(f'### SUBSET: {args.subset}')
        print(f'#################################################')
        if hasattr(args, 'src'):
            print(f'### SRC: {src_show}')
        if hasattr(args, 'dst'):
            print(f'### DST: {dst_show}')
        if hasattr(args, 'extras'):
            print(f'### ARGS: {args.extras}')
        print(f'#################################################')

####################################

class Paginator:

    def __init__(self):
        self.session = boto3.Session(
            aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'],
            aws_secret_access_key=os.environ['AWS_SECRET_ACCESS_KEY'],
        )
        self.s3 = self.session.client('s3')
        self.paginator = self.s3.get_paginator('list_objects_v2')
        self.bucket = 'tri-ml-sandbox-16011-us-west-2-datasets'

    def ls_seq(self, prefix):

        if prefix.startswith('s3://'):
            prefix = prefix[len(f's3://{self.bucket}')+1:]

        len_prefix = len(prefix)
        if not prefix.endswith('/'):
            len_prefix += 1

        tree = {}
        for page in self.paginator.paginate(Bucket=self.bucket, Prefix=prefix):
            for obj in page.get('Contents', []):

                key = obj['Key']
                folder = os.path.dirname(key[len_prefix:])
                folder, camera = os.path.dirname(folder), os.path.basename(folder)
                tarfile = os.path.basename(key)

                if camera not in tree:
                    tree[camera] = []
                tree[camera].append(tarfile)

        return tree

    def ls_all(self, prefix):

        if prefix.startswith('s3://'):
            prefix = prefix[len(f's3://{self.bucket}')+1:]

        len_prefix = len(prefix)
        if not prefix.endswith('/'):
            len_prefix += 1

        tree = {}
        for page in self.paginator.paginate(Bucket=self.bucket, Prefix=prefix):
            for obj in page.get('Contents', []):

                key = obj['Key']
                folder = os.path.dirname(key[len_prefix:])
                folder, camera = os.path.dirname(folder), os.path.basename(folder)
                tarfile = os.path.basename(key)

                if folder not in tree:
                    tree[folder] = dict()
                if camera not in tree[folder]:
                    tree[folder][camera] = []
                tree[folder][camera].append(tarfile)

        return tree

####################################

class Counters:

    def __init__(self, args, seqs):

        mp_manager = mp.Manager()
        self.counter = mp_manager.Value('counter', 0)
        self.suc = mp_manager.Value('suc', 0)
        self.err = mp_manager.Value('err', 0)
        self.don = mp_manager.Value('don', 0)
        self.prc = mp_manager.Value('prc', 0)
        self.lock = mp.Lock()

        self.n_seqs = len(seqs)
        self.n_procs = args.num_procs

    def click_don(self):
        with self.lock:
            self.don.value += 1
            self.prc.value += 1
            # if self.prc.value >= self.n_seqs:
            #     self.counter.value += 1

    def click_suc(self):
        with self.lock:
            self.suc.value += 1
            self.prc.value += 1
            # if self.prc.value >= self.n_seqs:
            #     self.counter.value += 1

    def click_err(self):
        with self.lock:
            self.err.value += 1
            self.prc.value += 1
            # if self.prc.value >= self.n_seqs:
            #     self.counter.value += 1

    def click_counter(self, seqs):
        if self.counter.value >= self.n_seqs:
            while self.counter.value < self.n_seqs + self.n_procs:
                time.sleep(1)
            return None

        with self.lock:
            value = self.counter.value
            self.counter.value += 1
            return value

    def display(self, args, progress, text, click_progress=False):
        with self.lock:
            description = f'| SED: {self.suc.value}/{self.err.value}/{self.don.value}'
            if args.subset is not None: description = f': {args.subset} {description}'
            progress.set_description(f'# {text} {description}')
            progress.n = self.prc.value
            progress.refresh()
            # if self.prc.value >= self.n_seqs:
            #     self.counter.value += 1

####################################