import random
import threading
from queue import Queue, Empty
from typing import Any, List, Sized, Dict, Protocol


DEBUG = False

def log_debug(msg):
    if DEBUG:
        print(f"[{threading.current_thread().name}] {msg}", flush=True)


class Dataset(Protocol):
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> Any: ...
    @staticmethod
    def collate(batch: List[Any]) -> Any: ...


class StepBasedSampler:
    def __init__(self, data_source: Sized, seed: int = 0,
                 rank: int = 0, world_size: int = 1,
                 global_batch_size: int = 1):
        self.data_source = data_source
        self.n = len(data_source)
        self.seed = seed
        self.rank = rank
        self.world_size = world_size
        self.global_batch_size = global_batch_size
        
        if global_batch_size % world_size != 0:
            raise ValueError(
                f"Global batch size {global_batch_size} must be divisible by world size {world_size}")
        
        self.local_batch_size = global_batch_size // world_size
    
    def get_indices_for_step(self, global_step: int) -> List[int]:
        combined_seed = self.seed * 10000 + global_step
        rng = random.Random(combined_seed)
        global_indices = rng.choices(range(self.n), k=self.global_batch_size)
        start_idx = self.rank * self.local_batch_size
        end_idx = start_idx + self.local_batch_size
        return global_indices[start_idx:end_idx]


class StepBasedDataLoader:
    def __init__(self,
                 dataset,
                 global_batch_size: int = 1,
                 seed: int = 0,
                 num_workers: int = 0,
                 rank: int = 0,
                 world_size: int = 1):
        self.dataset = dataset
        self.global_batch_size = global_batch_size
        self.seed = seed
        self.num_workers = max(0, num_workers)
        self.rank = rank
        self.world_size = world_size
        self.sampler = StepBasedSampler(
            dataset,
            seed=seed,
            rank=rank,
            world_size=world_size,
            global_batch_size=global_batch_size
        )
        
        self.stop_event = threading.Event()
        self.cache: Dict[int, List] = {}
        self.current_step = -1
        self.lock = threading.Lock()
        
        # Multi-process components
        self.task_queue: Queue = Queue()
        self.result_queue: Queue = Queue()
        self.workers = []
        self.result_thread = None
        
        if num_workers > 0:
            self._start_workers()
    
    def _start_workers(self):
        """Start worker processes"""
        log_debug(f"Starting {self.num_workers} workers")
        for i in range(self.num_workers):
            worker = threading.Thread(
                target=self._worker_main,
                args=(i, self.dataset, self.task_queue, self.result_queue, self.stop_event),
                daemon=True
            )
            worker.start()
            self.workers.append(worker)
            log_debug(f"Worker {i} started")
    
    @staticmethod
    def _worker_main(worker_id, dataset, task_queue: Queue, result_queue: Queue, stop_event):
        """Worker thread main loop"""
        log_debug(f"Worker {worker_id} running")
        try:
            while not stop_event.is_set():
                try:
                    task = task_queue.get(timeout=0.1)
                    if task is None:  # Shutdown signal
                        break
                    
                    step, position, index = task
                    sample = dataset[index]
                    result_queue.put((step, position, sample))
                except Empty:
                    continue
        except Exception as e:
            log_debug(f"Worker {worker_id} error: {str(e)}")
        finally:
            log_debug(f"Worker {worker_id} exiting")
    
    def _fetch_batch(self, step: int, dispatch: bool) -> List:
        """Fetch and cache a batch"""
        if step in self.cache:
            return self.cache[step]
        
        indices = self.sampler.get_indices_for_step(step)
        batch_size = len(indices)
        batch = [None] * batch_size
        
        if dispatch:
            # Dispatch tasks
            for position, index in enumerate(indices):
                self.task_queue.put((step, position, index))
            
            # Collect results
            count = 0
            while count < batch_size:
                s, position, sample = self.result_queue.get()
                if s == step:
                    batch[position] = sample
                    count += 1
        else:
            # Single-process loading
            for position, index in enumerate(indices):
                batch[position] = self.dataset[index]
        
        self.cache[step] = batch
        return batch
    
    def get_batch_at_step(self, step: int) -> Any:
        """Get batch for specified step with prefetch"""
        with self.lock:
            # Fetch current batch
            if step not in self.cache:
                batch = self._fetch_batch(step, False)
            else:
                batch = self.cache[step]
            
            # Prefetch next batch
            if self.num_workers > 0:
                self.current_step = step
                next_step = step + 1
                threading.Thread(
                    target=self._fetch_batch,
                    args=(next_step, True),
                    daemon=True
                ).start()
            
            # Cleanup old cache entries
            for key in list(self.cache.keys()):
                if key < step - 1:  # Keep only current and previous batch
                    del self.cache[key]
            
            return self.dataset.collate(batch)
    
    def shutdown(self):
        """Cleanup resources"""
        self.stop_event.set()
        for _ in range(self.num_workers):
            self.task_queue.put(None)  # Send shutdown signal


def test_dataloader():
    import torch
    import numpy as np
    import time

    # 模拟耗时数据集（增加唯一标识）
    class VerifiableDataset:
        def __init__(self, size=1000, seed=0):
            self.size = size
            rng = np.random.default_rng(seed)
            # 生成唯一可验证数据
            self.unique_ids = np.arange(size)
            rng.shuffle(self.unique_ids)
            self.data = self.unique_ids[:, None].astype(np.float32)

        def __len__(self):
            return self.size

        def __getitem__(self, index):
            time.sleep(0.001)  # 1ms处理时间
            return self.data[index], index  # 返回数据值和原始索引

        @staticmethod
        def collate(batch):
            data, indices = zip(*batch)
            return torch.tensor(np.array(data)), torch.tensor(np.array(indices))

    # 辅助函数：检查批次唯一性
    def check_batch_uniqueness(batch_indices):
        unique_indices = set(batch_indices.numpy().flatten())
        # print(f"Batch indices: {batch_indices}")
        return len(unique_indices) == len(batch_indices), len(unique_indices)

    print("===== 测试1: 零worker模式 =====")
    dataset = VerifiableDataset(size=1000, seed=1)
    dataloader = StepBasedDataLoader(
        dataset,
        global_batch_size=32,
        num_workers=0,
        rank=0,
        world_size=1,
        seed=1,
    )

    # 验证数据可重复性
    step0_data1, _ = dataloader.get_batch_at_step(0)
    step0_data2, _ = dataloader.get_batch_at_step(0)
    assert torch.allclose(step0_data1, step0_data2), "零worker模式数据不一致"
    print("零worker模式验证通过")

    print("\n===== 测试2: 边界条件 =====")
    # 测试最后一步加载
    last_step = len(dataset) // 32
    print(f"获取最后一步 (step={last_step})...")
    last_batch, last_indices = dataloader.get_batch_at_step(last_step)
    print(f"最后批次大小: {last_batch.shape[0]}")

    # 检查唯一性
    is_unique, unique_count = check_batch_uniqueness(last_indices)
    print(f"最后批次唯一索引: {unique_count}/{last_batch.shape[0]} (唯一)")

    # 测试超范围步骤
    try:
        dataloader.get_batch_at_step(last_step + 100)
        assert False, "未触发超范围步骤异常"
    except Exception as e:
        print(f"超范围步骤异常测试通过: {str(e)}")

    dataloader.shutdown()

    print("\n===== 测试3: 预取缓存有效性 =====")
    dataloader = StepBasedDataLoader(
        VerifiableDataset(size=1000, seed=1),
        global_batch_size=32,
        num_workers=4,
        rank=0,
        world_size=1,
        seed=1,
    )

    # 首次访问（无预取）
    start_time = time.time()
    step0_data, step0_indices = dataloader.get_batch_at_step(0)
    step0_time = time.time() - start_time
    print(f"步骤0: 耗时={step0_time:.4f}s (无预取)")

    time.sleep(0.1)  # 模拟GPU训练

    # 第二次访问（有预取）
    start_time = time.time()
    step1_data, step1_indices = dataloader.get_batch_at_step(1)
    step1_time = time.time() - start_time
    print(f"步骤1: 耗时={step1_time:.4f}s (有预取)")

    # 验证预取加速效果
    assert step1_time < step0_time * 0.7, "预取未显著减少加载时间"
    print(f"预取加速比: {step0_time/step1_time:.2f}x")

    # 验证预取缓存正确性
    is_unique, unique_count = check_batch_uniqueness(step1_indices)
    assert is_unique, f"步骤1批次索引重复: {unique_count}个唯一索引"
    print(f"步骤1批次验证: {unique_count}个唯一索引")

    # 测试预取干扰
    start_time = time.time()
    step10_data, step10_indices = dataloader.get_batch_at_step(10)
    step10_time = time.time() - start_time
    print(f"步骤10: 耗时={step10_time:.4f}s (无预取)")

    # 验证预取不影响随机访问
    assert step10_time > step1_time * 1.5, "预取干扰随机访问"

    # 再次访问步骤1
    start_time = time.time()
    step1_data_again, _ = dataloader.get_batch_at_step(1)

    # 验证缓存数据一致性
    assert torch.allclose(step1_data, step1_data_again), "缓存数据不一致"

    # 预取验证
    time.sleep(0.1)  # 模拟GPU训练
    start_time = time.time()
    step2_data, _ = dataloader.get_batch_at_step(2)
    step2_time = time.time() - start_time
    print(f"步骤2: 耗时={step2_time:.4f}s (有预取)")

    dataloader.shutdown()

    print("\n===== 测试4: DDP模式预取 =====")
    world_size = 2
    global_batch_size = 32

    def run_rank(rank):
        dataset = VerifiableDataset(size=1000, seed=3)
        dataloader = StepBasedDataLoader(
            dataset,
            global_batch_size=global_batch_size,
            num_workers=2,
            rank=rank,
            world_size=world_size
        )

        # 测试预取加速
        step0_start = time.time()
        dataloader.get_batch_at_step(0)
        step0_time = time.time() - step0_start

        time.sleep(0.1)  # 模拟GPU训练
        step1_start = time.time()
        dataloader.get_batch_at_step(1)
        step1_time = time.time() - step1_start

        print(f"Rank {rank}: 步骤0耗时={step0_time:.4f}s, 步骤1耗时={step1_time:.4f}s")

        dataloader.shutdown()
        return step0_time, step1_time

    # 运行多rank
    from concurrent.futures import ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=world_size) as executor:
        results = list(executor.map(run_rank, range(world_size)))

    # 验证所有rank的预取加速
    for rank, (step0_time, step1_time) in enumerate(results):
        assert step1_time < step0_time, f"Rank {rank} 预取未加速"
        print(f"Rank {rank} 预取加速比: {step0_time/step1_time:.2f}x")

    print("\n===== 所有测试通过 =====")


if __name__ == "__main__":
    print("===== 增强版数据加载器测试开始 =====")
    test_dataloader()
    print("===== 测试结束 =====")
