import functools
import os
import traceback

from imaginaire.utils import distributed, log  # NOTE(bvh): very slow for some reason?
from custom.utils.aws import aws_s3_cp, aws_s3_sync


# SageMaker always uses /tmp/a4d2 (28 TB scratch), ignoring config overrides.
_SM_LOCAL_ROOT = '/tmp/a4d2'
_IS_SAGEMAKER = os.getenv('SAGEMAKER', None) == 'enabled'
_DEFAULT_LOCAL_ROOT = _SM_LOCAL_ROOT if _IS_SAGEMAKER else '/any4d'


def get_local_root(config_job=None):
    """Return the local filesystem root, from config or default.
    On SageMaker, always returns /tmp/a4d2 regardless of config.
    :param config_job: JobConfig (or None). Reads local_root if available.
    """
    if _IS_SAGEMAKER:
        return _SM_LOCAL_ROOT
    if config_job is not None:
        lr = getattr(config_job, 'local_root', '')
        if lr:
            return lr
    return _DEFAULT_LOCAL_ROOT


def resilient_subroutine(max_failure_rate=0.20, grace_period=10, default=None):
    '''Decorator: catch exceptions and re-raise only if failure rate exceeds threshold.
    :param max_failure_rate: Fraction of calls allowed to fail (0.15 = 15%).
    :param grace_period: Minimum number of calls before enforcing the rate.
    :param default: Return value (or callable producing it) on failure.
    '''
    # TODO(bvh): unify this with the robustness logic in
    # custom/dataset/anydata_dataset.py:AnyDataset.__getitem__()
    
    def decorator(fn):
        fn._resilient_subroutine_calls = 0
        fn._resilient_subroutine_failures = 0

        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            fn._resilient_subroutine_calls += 1
            try:
                return fn(*args, **kwargs)
            except Exception:
                fn._resilient_subroutine_failures += 1
                n, f = fn._resilient_subroutine_calls, fn._resilient_subroutine_failures
                log.error(f'{fn.__name__} failed ({f}/{n} = {f/n:.0%}):\n'
                          f'{traceback.format_exc()}', rank0_only=False)
                if n >= grace_period and f / n > max_failure_rate:
                    raise RuntimeError(
                        f'{fn.__name__} failure rate too high: {f}/{n}') from None
                return default() if callable(default) else default

        return wrapper
    
    return decorator


def get_root_s3(path):
    """Returns s3 root path, which is different when loading pre-trained weights or fine-tuned checkpoints"""
    if '/sagemaker/' in path:
        return '/'.join(path.split('/')[:5])        # Fine-tuned (one folder deeper)
    else:
        return '/'.join(path.split('/')[:4])        # NVIDIA pre-trained weights


def get_checkpoint(path, local_root=''):
    """Get local checkpoint path for a single file, and if a s3 bucket is provided downloads it first.
    :param path: S3 or local path to the checkpoint file.
    :param local_root: Local filesystem root for caching S3 downloads. Empty = use default.
    """
    if path is None:
        return path
    if path.startswith('s3://'):
        root_s3 = get_root_s3(path)
        local_folder = path.replace(root_s3, local_root or _DEFAULT_LOCAL_ROOT)
        if distributed.is_local_rank0():
            if not os.path.exists(local_folder):
                aws_s3_cp(path, local_folder)
        distributed.barrier()
        return local_folder
    distributed.barrier()
    return path


def get_folder(path, local_root=''):
    """Get local checkpoint path for a folder, and if a s3 bucket is provided downloads it first.
    :param path: S3 or local path to the folder.
    :param local_root: Local filesystem root for caching S3 downloads. Empty = use default.
    """
    if path is None:
        return path
    if path.startswith('s3://'):
        root_s3 = get_root_s3(path)
        local_folder = path.replace(root_s3, local_root or _DEFAULT_LOCAL_ROOT)
        if distributed.is_local_rank0():
            if not os.path.exists(local_folder):
                aws_s3_sync(path, local_folder)
        distributed.barrier()
        return local_folder
    distributed.barrier()
    return path


def get_resume(path, local_root=''):
    """Get resume checkpoints (model/optim/scheduler/trainer), and if a s3 bucket is provided downloads each one first.
    :param path: S3 or local path to one of the checkpoint files (typically model/).
    :param local_root: Local filesystem root for caching S3 downloads. Empty = use default.
    """
    if path is None:
        return path
    if path.startswith('s3://'):
        root_s3 = get_root_s3(path)
        _root = local_root or _DEFAULT_LOCAL_ROOT
        if distributed.is_local_rank0():
            # Use model path as the base, or extra if using additional networks
            base = '/model/' if '/model/' in path else '/extra/'
            for folder in ['model', 'extra', 'optim', 'scheduler', 'trainer']:
                s3_folder = path.replace(base, f'/{folder}/') # Replace base with current folder
                local_folder = s3_folder.replace(root_s3, _root)
                if not os.path.exists(local_folder):
                    # NOTE(bvh): 'extra' (additional-network state) only exists if the model has
                    # extra_nets; tolerate its absence (robust) instead of crashing the resume.
                    # model/optim/scheduler/trainer stay strict -- a real resume needs them.
                    aws_s3_cp(s3_folder, local_folder, robust=(folder == 'extra'))
        distributed.barrier()
        return path.replace(root_s3, _root)
    distributed.barrier()
    return path