import random
from typing import Dict

from anydata.dataloaders.BaseWebbed import BaseWebbedDataset, pytorch_worker_info

from copy import deepcopy


def get_rank_worker():
    rank, _, worker, _ = pytorch_worker_info()
    return f'{rank}_{worker}'


class WebbedDataset(BaseWebbedDataset):
    def __init__(self, shuffle_samples=False, single_sample=None, subsample=None, **kwargs):
        super().__init__(buffer=True, subsample=subsample, **kwargs)
        self.super_sample: Dict[str, dict] = {}
        self.available: Dict[str, list] = {}

        self.shuffle = shuffle_samples
        self.single_sample = single_sample

    def __len__(self):
        # NOTE(bvh): 100M is large but bounded (2^31 would hang DistributedSampler).
        # Keep this value consistent between WebbedDataset (AnyData) and RankDataset (Any4D / custom).
        if self.shard_type_name == 'resampled':
            return 100_000_000
        else:
            # num_tarfiles already reflects subsample filtering from BaseWebbed
            return self.num_tarfiles

    def get_rank_worker(self):
        rank, _, worker, _ = pytorch_worker_info()
        return f'{rank}_{worker}'

    def create_sample(self, rank_worker):
        # Invalidate any prior super_sample cache to prevent confusing error messages
        self.super_sample[rank_worker] = {}

        # TODO(bvh): if a crash ever occurs inside here (and not inside the already
        # robustified parts of tariterators), we may need similar logic context
        # attaching logic, but this is quite heavy so skipping for now.
        super_sample = super().__getitem__(0)

        keys = super_sample['rgb'].keys()
        time = sorted(list(set([k[0] for k in keys])))
        total, length = len(time), len(self.context)
        if total <= length:
            # NOTE(bvh): targeted message for the config-mismatch case so the
            # generic warn-and-retry in AnyDataset surfaces a clear cause
            # instead of `random.choice([])` IndexError below.
            raise ValueError(
                f"Webbed super_sample too short: have {total} unique timesteps "
                f"but length requires {length + 1}. Likely cause: each "
                f"super_sample may be constrained by `multi_tarfiles`")
        available = list(range(total - length))

        # Update stride to include webdataset stride
        super_sample['metadata']['stride'] = {
            key: val * self.context_stride for key, val in super_sample['metadata']['stride'].items()}

        if self.single_sample == 'first':
            available = [0] # Get only first
        elif self.single_sample == 'random':
            available = [random.choice(available)] # Pick random one
        if self.shuffle:
            random.shuffle(available) # Shuffle available if requested

        self.super_sample[rank_worker] = super_sample
        self.available[rank_worker] = available

    def build_sample(self, rank_worker):
        super_sample = self.super_sample[rank_worker]
        available = self.available[rank_worker]

        start = available[0]
        self.available[rank_worker] = available[1:]
        raw_stride = super_sample['metadata']['stride']
        # stride can be: int (old format), or dict {cam_idx: value} (new format)
        if isinstance(raw_stride, dict):
            stride_val = raw_stride[0] if 0 in raw_stride else next(iter(raw_stride.values()))
            stride = next(iter(stride_val.values())) if isinstance(stride_val, dict) else stride_val
        else:
            stride = raw_stride
        
        start = start * stride // self.context_stride
        finish = start + stride * (len(self.context) + 1) // self.context_stride

        sample = {'metadata': super_sample['metadata']}
        for label in super_sample.keys():
            if label == 'metadata':
                continue
            sample[label] = {(key[0] - start, key[1]): val 
                for key, val in super_sample[label].items()
                if start <= key[0] < finish and (key[0] - start) % (stride) == 0}

        # Offset timestep and store timestep_start
        # Use first available camera (not hardcoded 0) since some datasets
        # may not have camera 0 after core_aux_loaded filtering
        first_cam = min(k[1] for k in sample['timestep'])
        timestep_start = deepcopy(sample['timestep'][(0, first_cam)])
        for key, val in sample['timestep'].items():
            sample['timestep'][key] = val - timestep_start
        sample['metadata']['timestep_start'] = timestep_start

        return sample

    def _do_getitem(self, rank_worker):
        # Create new super sample if there is none or we ran out of availability
        if rank_worker not in self.available.keys() or len(self.available[rank_worker]) == 0:
            self.create_sample(rank_worker)

        # Fetch super sample and availablility
        sample = self.build_sample(rank_worker)
        # sample1 = deepcopy(sample)  # DEBUG(bvh)

        sample2 = self.post_process_sample(sample, webbed=True)
        return sample2

    def __getitem__(self, idx=None):
        # Get rank and worker for this instance
        rank_worker = self.get_rank_worker()  # e.g. '1_2'
        try:
            return self._do_getitem(rank_worker)
        
        except Exception as exc:
            # Attach per worker diagnostic state so upstream users
            # (e.g. custom AnyDataset's _log_error_inspection) can surface
            # which shard the failure came from
            if not hasattr(exc, 'webbed_context'):

                ss = self.super_sample.get(rank_worker)  # may be None
                meta = ss.get('metadata', {}) if isinstance(ss, dict) else {}

                exc.webbed_context = {
                    'rank_worker': rank_worker,
                    'dset_name': meta.get('name') or (meta.get('info', {}) or {}).get('name'),
                    'raw_id': meta.get('raw_id'),
                    'path': meta.get('path'),
                    'super_sample': ss,
                    'available_remaining': self.available.get(rank_worker, []),
                }

            # NOTE(bvh): force next call through create_sample; otherwise one
            # bad super_sample cascades ~N sub-clip failures and trips AnyDataset's
            # error threshold.
            self.available[rank_worker] = []

            # Always raise errors to the caller, since robustness is not our responsibility here,
            # only adding information for upstream wrappers to understand and handle.
            raise