import h5py
import remfile
import boto3
from botocore import UNSIGNED
from botocore.config import Config
from pynwb import NWBHDF5IO

S3_BUCKET = "aind-open-data"

def _get_s3_client():
    return boto3.client("s3", config=Config(signature_version=UNSIGNED))


def _parse_s3_path(s3_path, default_bucket=S3_BUCKET):
    """Parse an s3:// URL or plain folder name into (bucket, key)."""
    if s3_path.startswith("s3://"):
        s3_path = s3_path[len("s3://"):]
        bucket, _, key = s3_path.partition("/")
        return bucket, key
    return default_bucket, s3_path

def _list_all_keys(prefix, bucket=S3_BUCKET):
    """Recursively list all S3 object keys under a prefix."""
    prefix = prefix.rstrip("/") + "/"
    paginator = _get_s3_client().get_paginator("list_objects_v2")
    keys = []
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
        for obj in page.get("Contents", []):
            keys.append(obj["Key"])
    return keys


def _find_nwb_in_keys(keys):
    """
    Search a list of S3 keys for NWB files.
    Returns (nwb_key_or_prefix, is_zarr), or (None, None) if not found.
    """
    zarr_prefixes = set()
    hdf5_keys = []

    for key in keys:
        if ".nwb.zarr/" in key:
            zarr_prefix = key[: key.index(".nwb.zarr/") + len(".nwb.zarr")]
            zarr_prefixes.add(zarr_prefix)
        elif ".nwb/" in key:
            zarr_prefix = key[: key.index(".nwb/") + len(".nwb")]
            zarr_prefixes.add(zarr_prefix)
        elif key.endswith(".nwb"):
            hdf5_keys.append(key)

    if zarr_prefixes:
        return next(iter(zarr_prefixes)), True
    if hdf5_keys:
        return hdf5_keys[0], False
    return None, None

def list_s3_folder(folder_prefix, bucket=S3_BUCKET):
    """
    List immediate children (files and subfolders) of a folder prefix in S3.

    Parameters
    ----------
    folder_prefix : str
        S3 folder prefix (e.g., "multiplane-ophys_832700_2026-01-29_11-18-09")
    bucket : str, optional
        S3 bucket name. Defaults to "aind-open-data"

    Returns
    -------
    list of str
        Keys of files and folder prefixes directly under the given prefix
    """
    prefix = folder_prefix.rstrip("/") + "/"
    s3_client = _get_s3_client()
    paginator = s3_client.get_paginator("list_objects_v2")

    results = []
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
        for obj in page.get("Contents", []):
            results.append(obj["Key"])
        for common_prefix in page.get("CommonPrefixes", []):
            results.append(common_prefix["Prefix"])

    return results

def stream_s3_file(object_key, bucket=S3_BUCKET):
    """
    Stream a file from S3 as a file-like object.

    Parameters
    ----------
    object_key : str
        S3 object key/path (e.g., "path/to/file.nwb")
    bucket : str, optional
        S3 bucket name. Defaults to "aind-open-data"

    Returns
    -------
    StreamingBody
        A file-like object for reading the S3 object
    """
    response = _get_s3_client().get_object(Bucket=bucket, Key=object_key)
    return response["Body"]

def open_nwb_from_s3(s3_path, bucket=S3_BUCKET):
    """
    Find and open an NWB file within an S3 folder.

    Recursively searches the given S3 path for an NWB file, determines whether
    it is HDF5 or Zarr, and opens it with the appropriate library.

    Parameters
    ----------
    s3_path : str
        Either an s3:// URL (e.g., "s3://aind-open-data/session_folder") or
        a plain folder name (e.g., "multiplane-ophys_832700_2026-01-29_11-18-09").
    bucket : str, optional
        S3 bucket to use when s3_path is a plain folder name. Defaults to "aind-open-data".

    Returns
    -------
    pynwb.NWBFile
        The opened NWB file object.
    """
    bucket, key = _parse_s3_path(s3_path, default_bucket=bucket)

    print(f"Searching for NWB file in s3://{bucket}/{key} ...")
    all_keys = _list_all_keys(key, bucket=bucket)
    nwb_key, is_zarr = _find_nwb_in_keys(all_keys)

    if nwb_key is None:
        raise FileNotFoundError(f"No NWB file found under s3://{bucket}/{key}")

    if is_zarr:
        try:
            from hdmf_zarr import NWBZarrIO
        except ImportError as e:
            raise ImportError(
                f"Zarr NWB found at s3://{bucket}/{nwb_key} but hdmf_zarr is unavailable: {e}"
            )
        print(f"Opening Zarr NWB: s3://{bucket}/{nwb_key}")
        io = NWBZarrIO(
            path=f"s3://{bucket}/{nwb_key}",
            mode="r",
            load_namespaces=True,
            storage_options={"anon": True},
        )
    else:
        print(f"Opening HDF5 NWB: s3://{bucket}/{nwb_key}")
        url = f"https://s3.amazonaws.com/{bucket}/{nwb_key}"
        rem = remfile.File(url)
        h5_file = h5py.File(rem, "r")
        io = NWBHDF5IO(file=h5_file, mode="r", load_namespaces=True)

    return io.read()