# 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 externals/AnyData/anydata/sync/sync_utils.py (shared retry/robust logic).
# NOTE(bvh): keep aws_s3_cp/sync/rm behavior consistent with sync_utils.py.

import subprocess
import time
import traceback

from rich.console import Console

_console = Console(stderr=True)
_RETRY_COUNT = 3
_RETRY_DELAY = 2  # 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. Previously each
    # frame was a separate write, producing the "_do_sync + tqdm bar" mash-up
    # seen in CloudWatch / wandb logs.
    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)


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

    Always retries up to _RETRY_COUNT times on failure.
    robust=False (default): raise RuntimeError on final failure (unless check=False).
      Use this for operations essential to training (e.g. loading initial weights).
    robust=True: on final failure, print a loud warning + stack trace but return False.
      Use this for non-essential operations (e.g. uploading visuals/checkpoints).
    """
    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(f'{err_msg}\nContinuing training.')
        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 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(f'{err_msg}\nContinuing training.')
        return False
    raise RuntimeError(err_msg)


def aws_s3_rm(path, recursive=False, quiet=True, extras='', robust=False):
    """Remove object(s) 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(f'{err_msg}\nContinuing training.')
        return False
    raise RuntimeError(err_msg)


def aws_s3_ls(s3_path, suffix=None):
    """List immediate children of an S3 path (files and folder names). Raises on failure."""
    result = subprocess.run(
        f'aws s3 ls {s3_path}/', shell=True, capture_output=True, text=True
    )
    if result.returncode != 0:
        raise RuntimeError(f'aws s3 ls failed: {s3_path}\n{result.stderr.strip()}')
    data = result.stdout.replace('/\n', ' ').replace('\n', ' ').replace('  ', ' ').split(' ')
    data = [d for d in data if d and d not in ['PRE']]
    if suffix is not None:
        data = [d for d in data if d.endswith(suffix)]
    return data


def aws_s3_ls_folders(s3_path):
    """List folder names (PRE entries) at an S3 path. Raises on failure."""
    result = subprocess.run(
        f'aws s3 ls {s3_path}/', shell=True, capture_output=True, text=True
    )
    if result.returncode != 0:
        raise RuntimeError(f'aws s3 ls failed: {s3_path}\n{result.stderr.strip()}')
    data = result.stdout.replace('/\n', ' ').replace('\n', ' ').replace('  ', ' ').split(' ')
    return [data[i + 1] for i in range(len(data) - 1) if data[i] == 'PRE']


def aws_s3_ls_recursive(s3_path, suffix=None):
    """Recursively list all objects under an S3 path. Raises on failure."""
    result = subprocess.run(
        f'aws s3 ls {s3_path}/ --recursive', shell=True, capture_output=True, text=True
    )
    if result.returncode != 0:
        raise RuntimeError(f'aws s3 ls --recursive failed: {s3_path}\n{result.stderr.strip()}')
    data = result.stdout.replace('/\n', ' ').replace('\n', ' ').replace('  ', ' ').split(' ')
    data = [d for d in data if d]
    if suffix is not None:
        if isinstance(suffix, list):
            data = [d for d in data if any(d.endswith(s) for s in suffix)]
        else:
            data = [d for d in data if d.endswith(suffix)]
    return data