# import webdataset.gopen as wds_gopen
import importlib
import os
import sys
import threading
import time
wds_gopen = importlib.import_module("webdataset.gopen")

from anydata.utils.retry import retry_state

# Where aws --debug output gets dumped on retry. Override via env var if needed.
_AWS_DEBUG_DIR = os.environ.get('ANYDATA_AWS_DEBUG_DIR', '/tmp/aws_debug')


def register_s3_gopen():
    """Register an s3:// reader for external webdataset gopen."""
    if 's3' in wds_gopen.gopen_schemes:
        return

    # NOTE(bvh): explicit --cli-read-timeout / --cli-connect-timeout to bound hangs.
    # Without these, it could freeze the per-camera reader thread inside webdataset:tar_file_iterator.
    _timeouts = "--cli-read-timeout 60 --cli-connect-timeout 30"

    def gopen_s3(url, mode="rb", bufsize=8192):
        # On retry attempts add --debug so aws CLI dumps full HTTP/auth trace.
        # Redirect that (multi-MB) stderr to a per-call file to avoid drowning
        # the dataloader's own stderr.
        attempt = getattr(retry_state, 'attempt', 1)
        if attempt > 1:
            os.makedirs(_AWS_DEBUG_DIR, exist_ok=True)
            log_path = (
                f"{_AWS_DEBUG_DIR}/pid{os.getpid()}_tid{threading.get_ident()}"
                f"_a{attempt}_{int(time.time())}.log"
            )
            debug_flag = "--debug"
            redirect = f"2>>'{log_path}'"
            print(
                f"[gopen_s3] retry attempt={attempt} url={url} aws --debug -> {log_path}",
                file=sys.stderr, flush=True,
            )
        else:
            debug_flag = ""
            redirect = ""

        if mode[0] == "r":
            cmd = f"aws {debug_flag} {_timeouts} s3 cp '{url}' - {redirect}"
            return wds_gopen.Pipe(
                cmd,
                mode=mode,
                shell=True,
                bufsize=bufsize,
                ignore_status=[141, 23],
            )
        if mode[0] == "w":
            cmd = f"aws {debug_flag} {_timeouts} s3 cp - '{url}' {redirect}"
            return wds_gopen.Pipe(
                cmd,
                mode=mode,
                shell=True,
                bufsize=bufsize,
                ignore_status=[141, 26],
            )
        raise ValueError(f"{mode}: unknown mode")

    wds_gopen.gopen_schemes['s3'] = gopen_s3