"""
Distributed Training System for NeuroLLM

Implements decentralized training where:
1. Nodes contribute compute for forward/backward passes
2. Gradients are aggregated via gossip protocol
3. Training rewards are distributed in NEURO tokens
4. Model checkpoints are shared across the network

Key Components:
- GradientAggregator: Collects and averages gradients from peers
- TrainingCoordinator: Orchestrates distributed training
- DataContributor: Handles federated dataset management
- RewardCalculator: Computes NEURO rewards for contributions

Training Flow:
1. Coordinator broadcasts current model state hash
2. Nodes with matching state participate in training round
3. Each node processes local data batch
4. Gradients are compressed and gossiped
5. Aggregated gradients are applied
6. New checkpoint is created and distributed
7. NEURO rewards are calculated and distributed
"""

import torch
import torch.nn as nn
import threading
import time
import hashlib
from concurrent.futures import ThreadPoolExecutor
import logging
import json
import io
import zlib
import base64
import os
import requests
from typing import Dict, List, Optional, Tuple, Any, Callable
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict

# Import economics constants for consistency
from neuroshard.core.economics.constants import (
    TRAINING_REWARD_PER_BATCH,
    DATA_REWARD_PER_SAMPLE
)

logger = logging.getLogger(__name__)


class TrainingState(Enum):
    """State of the training coordinator."""
    IDLE = "idle"
    COLLECTING = "collecting"      # Collecting gradients from peers
    AGGREGATING = "aggregating"    # Aggregating gradients
    APPLYING = "applying"          # Applying updates
    CHECKPOINTING = "checkpointing"  # Creating checkpoint


@dataclass
class GradientContribution:
    """A gradient contribution from a node."""
    node_id: str
    round_id: int
    layer_gradients: Dict[str, bytes]  # layer_name -> compressed gradient
    batch_size: int
    loss: float
    timestamp: float
    signature: str  # Proof of work


@dataclass
class TrainingRound:
    """A single training round."""
    round_id: int
    started_at: float
    model_hash: str
    
    # Contributions
    contributions: Dict[str, GradientContribution] = field(default_factory=dict)
    min_contributions: int = 3
    max_contributions: int = 100
    
    # Results
    aggregated_gradients: Optional[Dict[str, torch.Tensor]] = None
    total_batch_size: int = 0
    avg_loss: float = 0.0
    
    # State
    completed: bool = False
    applied: bool = False


@dataclass
class TrainingReward:
    """Reward for training contribution."""
    node_id: str
    round_id: int
    compute_reward: float      # For compute contribution
    data_reward: float         # For data contribution
    quality_bonus: float       # For high-quality gradients
    total_neuro: float


class GradientCompressor:
    """
    Compresses gradients for efficient network transmission.
    
    Uses a combination of:
    1. Top-K sparsification (keep only largest gradients)
    2. Quantization (reduce precision)
    3. zlib compression
    """
    
    def __init__(self, top_k_ratio: float = 0.1, bits: int = 8):
        self.top_k_ratio = top_k_ratio
        self.bits = bits
    
    def compress(self, gradient: torch.Tensor) -> bytes:
        """Compress a gradient tensor."""
        # CRITICAL: Move to CPU first for MPS/CUDA compatibility
        gradient = gradient.detach().cpu()
        
        # Flatten
        flat = gradient.flatten()
        
        # Top-K sparsification
        k = max(1, int(len(flat) * self.top_k_ratio))
        values, indices = torch.topk(flat.abs(), k)
        
        # Get actual values (with signs)
        sparse_values = flat[indices]
        
        # Quantize to specified bits
        max_val = sparse_values.abs().max()
        if max_val > 0:
            scale = (2 ** (self.bits - 1) - 1) / max_val
            quantized = (sparse_values * scale).round().to(torch.int8)
        else:
            quantized = torch.zeros(k, dtype=torch.int8)
            scale = 1.0
        
        # Pack into bytes (tensors already on CPU)
        data = {
            "shape": list(gradient.shape),
            "k": k,
            "indices": base64.b64encode(indices.numpy().tobytes()).decode('ascii'),
            "values": base64.b64encode(quantized.numpy().tobytes()).decode('ascii'),
            "scale": float(scale),
            "dtype": str(gradient.dtype),
        }
        
        # Serialize and compress
        json_data = json.dumps(data).encode()
        return zlib.compress(json_data)
    
    def decompress(self, data: bytes, device: str = "cpu") -> torch.Tensor:
        """Decompress a gradient tensor."""
        # Decompress and deserialize
        json_data = zlib.decompress(data)
        packed = json.loads(json_data)
        
        # Unpack
        shape = packed["shape"]
        k = packed["k"]
        indices = torch.frombuffer(
            bytearray(base64.b64decode(packed["indices"])), 
            dtype=torch.int64
        ).clone().to(device)
        values = torch.frombuffer(
            bytearray(base64.b64decode(packed["values"])), 
            dtype=torch.int8
        ).float().clone().to(device)
        scale = packed["scale"]
        
        # Dequantize
        values = values / scale
        
        # Reconstruct sparse tensor
        flat = torch.zeros(torch.prod(torch.tensor(shape)), device=device)
        flat[indices] = values
        
        return flat.view(*shape)


class GradientAggregator:
    """
    Aggregates gradients from multiple nodes.
    
    Supports:
    - Simple averaging
    - Weighted averaging (by batch size)
    - Robust aggregation (median, trimmed mean)
    """
    
    def __init__(self, method: str = "weighted_mean"):
        self.method = method
        self.compressor = GradientCompressor()
    
    def aggregate(
        self, 
        contributions: List[GradientContribution],
        layer_names: List[str]
    ) -> Dict[str, torch.Tensor]:
        """
        Aggregate gradients from multiple contributions.
        
        Args:
            contributions: List of gradient contributions
            layer_names: Names of layers to aggregate
            
        Returns:
            Aggregated gradients per layer
        """
        if not contributions:
            return {}
        
        aggregated = {}
        total_batch_size = sum(c.batch_size for c in contributions)
        
        for layer_name in layer_names:
            # Collect gradients for this layer
            gradients = []
            weights = []
            
            for contrib in contributions:
                if layer_name in contrib.layer_gradients:
                    grad = self.compressor.decompress(contrib.layer_gradients[layer_name])
                    gradients.append(grad)
                    weights.append(contrib.batch_size)
            
            if not gradients:
                continue
            
            # Stack gradients
            stacked = torch.stack(gradients)
            
            if self.method == "mean":
                aggregated[layer_name] = stacked.mean(dim=0)
                
            elif self.method == "weighted_mean":
                weights_tensor = torch.tensor(weights, dtype=torch.float32)
                weights_tensor = weights_tensor / weights_tensor.sum()
                aggregated[layer_name] = (stacked * weights_tensor.view(-1, *([1] * (stacked.dim() - 1)))).sum(dim=0)
                
            elif self.method == "median":
                aggregated[layer_name] = stacked.median(dim=0)[0]
                
            elif self.method == "trimmed_mean":
                # Remove top and bottom 10%
                k = max(1, len(gradients) // 10)
                sorted_grads = stacked.sort(dim=0)[0]
                aggregated[layer_name] = sorted_grads[k:-k].mean(dim=0) if k < len(gradients) // 2 else stacked.mean(dim=0)
        
        return aggregated


class TrainingCoordinator:
    """
    Coordinates distributed training across the network.
    
    Responsibilities:
    1. Initiate training rounds
    2. Collect gradient contributions
    3. Aggregate and apply updates
    4. Distribute rewards
    5. Manage checkpoints
    
    NOTE: This class is LEGACY and not currently used in production.
    The active reward path uses economics.py constants via ledger.py
    """
    
    # Configuration
    ROUND_DURATION_SECONDS = 60
    MIN_CONTRIBUTIONS = 3
    GRADIENT_CLIP_NORM = 1.0
    
    # Reward rates (using economics.py constants for consistency)
    # NOTE: LEGACY - These are kept for backwards compatibility but not actively used
    # Import at class level to match economics.py values
    from neuroshard.core.economics.constants import TRAINING_REWARD_PER_BATCH as COMPUTE_REWARD_PER_BATCH
    from neuroshard.core.economics.constants import DATA_REWARD_PER_SAMPLE
    QUALITY_BONUS_MULTIPLIER = 1.5
    
    def __init__(
        self,
        model: nn.Module,
        optimizer: torch.optim.Optimizer,
        node_id: str,
        ledger_manager = None,
        on_round_complete: Optional[Callable] = None
    ):
        self.model = model
        self.optimizer = optimizer
        self.node_id = node_id
        self.ledger = ledger_manager
        self.on_round_complete = on_round_complete
        
        # State
        self.state = TrainingState.IDLE
        self.current_round: Optional[TrainingRound] = None
        self.round_history: List[TrainingRound] = []
        self.global_step = 0
        
        # Components
        self.aggregator = GradientAggregator()
        self.compressor = GradientCompressor()
        
        # Threading
        self.lock = threading.Lock()
        self.running = False
        
        # Stats
        self.total_rounds = 0
        self.total_contributions = 0
        self.total_neuro_distributed = 0.0
        
        logger.info(f"TrainingCoordinator initialized for node {node_id}")
    
    def start(self):
        """Start the training coordinator."""
        self.running = True
        threading.Thread(target=self._training_loop, daemon=True).start()
        logger.info("Training coordinator started")
    
    def stop(self):
        """Stop the training coordinator."""
        self.running = False
    
    def _training_loop(self):
        """Main training loop."""
        while self.running:
            try:
                if self.state == TrainingState.IDLE:
                    # Start new round
                    self._start_round()
                    
                elif self.state == TrainingState.COLLECTING:
                    # Check if round should complete
                    if self._should_complete_round():
                        self._complete_round()
                        
                time.sleep(1)
                
            except Exception as e:
                logger.error(f"Training loop error: {e}")
                time.sleep(5)
    
    def _get_model_hash(self) -> str:
        """Get hash of current model state."""
        state_dict = self.model.state_dict()
        hasher = hashlib.sha256()
        
        for name, param in sorted(state_dict.items()):
            hasher.update(name.encode())
            hasher.update(param.cpu().numpy().tobytes()[:1000])  # Sample for speed
        
        return hasher.hexdigest()[:16]
    
    def _start_round(self):
        """Start a new training round."""
        with self.lock:
            self.total_rounds += 1
            
            self.current_round = TrainingRound(
                round_id=self.total_rounds,
                started_at=time.time(),
                model_hash=self._get_model_hash(),
                min_contributions=self.MIN_CONTRIBUTIONS,
            )
            
            self.state = TrainingState.COLLECTING
            
        logger.info(f"Started training round {self.total_rounds}")
    
    def _should_complete_round(self) -> bool:
        """Check if current round should complete."""
        if not self.current_round:
            return False
        
        # Time limit
        elapsed = time.time() - self.current_round.started_at
        if elapsed >= self.ROUND_DURATION_SECONDS:
            return True
        
        # Max contributions
        if len(self.current_round.contributions) >= self.current_round.max_contributions:
            return True
        
        return False
    
    def _complete_round(self):
        """Complete the current training round."""
        if not self.current_round:
            return
        
        with self.lock:
            round_data = self.current_round
            
            # Check minimum contributions
            if len(round_data.contributions) < round_data.min_contributions:
                logger.warning(f"Round {round_data.round_id} failed: insufficient contributions "
                             f"({len(round_data.contributions)}/{round_data.min_contributions})")
                self.state = TrainingState.IDLE
                self.current_round = None
                return
            
            self.state = TrainingState.AGGREGATING
        
        logger.info(f"Completing round {round_data.round_id} with {len(round_data.contributions)} contributions")
        
        # Aggregate gradients
        layer_names = [name for name, _ in self.model.named_parameters()]
        aggregated = self.aggregator.aggregate(
            list(round_data.contributions.values()),
            layer_names
        )
        
        round_data.aggregated_gradients = aggregated
        round_data.total_batch_size = sum(c.batch_size for c in round_data.contributions.values())
        round_data.avg_loss = sum(c.loss for c in round_data.contributions.values()) / len(round_data.contributions)
        
        # Apply gradients
        with self.lock:
            self.state = TrainingState.APPLYING
        
        self._apply_gradients(aggregated)
        round_data.applied = True
        
        # Calculate and distribute rewards
        rewards = self._calculate_rewards(round_data)
        self._distribute_rewards(rewards)
        
        # Checkpoint
        with self.lock:
            self.state = TrainingState.CHECKPOINTING
        
        self._create_checkpoint(round_data)
        
        # Complete
        round_data.completed = True
        self.round_history.append(round_data)
        
        # Keep only last 100 rounds
        if len(self.round_history) > 100:
            self.round_history = self.round_history[-100:]
        
        # Callback
        if self.on_round_complete:
            self.on_round_complete(round_data)
        
        # Reset
        with self.lock:
            self.current_round = None
            self.state = TrainingState.IDLE
            self.global_step += 1
        
        logger.info(f"Round {round_data.round_id} complete: loss={round_data.avg_loss:.4f}, "
                   f"batch_size={round_data.total_batch_size}")
    
    def _apply_gradients(self, gradients: Dict[str, torch.Tensor]):
        """Apply aggregated gradients to model."""
        self.optimizer.zero_grad()
        
        for name, param in self.model.named_parameters():
            if name in gradients:
                if param.grad is None:
                    param.grad = gradients[name].to(param.device)
                else:
                    param.grad.copy_(gradients[name])
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.GRADIENT_CLIP_NORM)
        
        # Apply
        self.optimizer.step()
    
    def _calculate_rewards(self, round_data: TrainingRound) -> List[TrainingReward]:
        """Calculate NEURO rewards for contributions."""
        rewards = []
        
        # Calculate average loss for quality comparison
        avg_loss = round_data.avg_loss
        
        for node_id, contrib in round_data.contributions.items():
            # Base compute reward
            compute_reward = contrib.batch_size * self.COMPUTE_REWARD_PER_BATCH
            
            # Data contribution reward
            data_reward = contrib.batch_size * self.DATA_REWARD_PER_SAMPLE
            
            # Quality bonus (lower loss = better)
            quality_bonus = 0.0
            if contrib.loss < avg_loss:
                quality_bonus = compute_reward * (self.QUALITY_BONUS_MULTIPLIER - 1)
            
            total = compute_reward + data_reward + quality_bonus
            
            rewards.append(TrainingReward(
                node_id=node_id,
                round_id=round_data.round_id,
                compute_reward=compute_reward,
                data_reward=data_reward,
                quality_bonus=quality_bonus,
                total_neuro=total
            ))
            
            self.total_neuro_distributed += total
        
        return rewards
    
    def _distribute_rewards(self, rewards: List[TrainingReward]):
        """Distribute NEURO rewards to contributors using PoNW proofs."""
        if not self.ledger:
            logger.debug("No ledger available for reward distribution")
            return
        
        for reward in rewards:
            try:
                from neuroshard.core.economics.ledger import PoNWProof, ProofType
                import time
                
                # Create a proper training PoNW proof
                proof = PoNWProof(
                    node_id=reward.node_id,
                    proof_type=ProofType.TRAINING.value,
                    timestamp=time.time(),
                    nonce=f"train_{reward.round_id}_{reward.node_id[:16]}",
                    training_batches=int(reward.compute_reward / self.COMPUTE_REWARD_PER_BATCH),
                    data_samples=int(reward.data_reward / self.DATA_REWARD_PER_SAMPLE),
                    signature=f"training_reward_{reward.round_id}_{reward.node_id}"
                )
                
                # Process through the ledger (handles deduplication, stats, etc.)
                success, amount, msg = self.ledger.process_proof(proof)
                
                if success:
                    logger.info(f"Reward: {reward.node_id[:8]}... earned {amount:.6f} NEURO "
                               f"(compute={reward.compute_reward:.6f}, data={reward.data_reward:.6f}, "
                               f"quality={reward.quality_bonus:.6f})")
                else:
                    logger.debug(f"Training reward not processed: {msg}")
                            
            except Exception as e:
                logger.error(f"Failed to distribute reward to {reward.node_id}: {e}")
    
    def _create_checkpoint(self, round_data: TrainingRound):
        """Create a checkpoint after training round."""
        checkpoint_path = f"checkpoints/neuro_llm_round_{round_data.round_id}.pt"
        
        try:
            import os
            os.makedirs("checkpoints", exist_ok=True)
            
            torch.save({
                "model_state_dict": self.model.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
                "round_id": round_data.round_id,
                "global_step": self.global_step,
                "model_hash": self._get_model_hash(),
                "avg_loss": round_data.avg_loss,
                "total_batch_size": round_data.total_batch_size,
                "timestamp": time.time(),
            }, checkpoint_path)
            
            logger.info(f"Checkpoint saved: {checkpoint_path}")
            
        except Exception as e:
            logger.error(f"Failed to save checkpoint: {e}")
    
    def submit_contribution(self, contribution: GradientContribution) -> bool:
        """
        Submit a gradient contribution for the current round.
        
        Called by peers when they have computed gradients.
        """
        with self.lock:
            if self.state != TrainingState.COLLECTING:
                return False
            
            if not self.current_round:
                return False
            
            # Verify model hash matches
            # In production, this would be more sophisticated
            
            # Add contribution
            self.current_round.contributions[contribution.node_id] = contribution
            self.total_contributions += 1
            
        logger.debug(f"Received contribution from {contribution.node_id[:8]}... "
                    f"(batch_size={contribution.batch_size}, loss={contribution.loss:.4f})")
        
        return True
    
    def compute_local_gradients(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor
    ) -> GradientContribution:
        """
        Compute gradients on local data.
        
        Call this to participate in training.
        """
        self.model.train()
        
        # Forward pass
        outputs = self.model(input_ids=input_ids, labels=labels)
        loss = outputs["loss"]
        
        # Backward pass
        loss.backward()
        
        # Collect and compress gradients
        layer_gradients = {}
        for name, param in self.model.named_parameters():
            if param.grad is not None:
                layer_gradients[name] = self.compressor.compress(param.grad)
        
        # Clear gradients (they're saved in contribution)
        self.optimizer.zero_grad()
        
        # Create contribution
        contribution = GradientContribution(
            node_id=self.node_id,
            round_id=self.current_round.round_id if self.current_round else 0,
            layer_gradients=layer_gradients,
            batch_size=input_ids.shape[0],
            loss=loss.item(),
            timestamp=time.time(),
            signature=self._sign_contribution(layer_gradients)
        )
        
        return contribution
    
    def _sign_contribution(self, gradients: Dict[str, bytes]) -> str:
        """Sign a contribution for verification."""
        hasher = hashlib.sha256()
        hasher.update(self.node_id.encode())
        hasher.update(str(time.time()).encode())
        for name, data in sorted(gradients.items()):
            hasher.update(name.encode())
            hasher.update(data[:100])  # Sample for speed
        return hasher.hexdigest()
    
    def get_status(self) -> Dict[str, Any]:
        """Get coordinator status."""
        return {
            "state": self.state.value,
            "global_step": self.global_step,
            "total_rounds": self.total_rounds,
            "total_contributions": self.total_contributions,
            "total_neuro_distributed": self.total_neuro_distributed,
            "current_round": {
                "round_id": self.current_round.round_id,
                "contributions": len(self.current_round.contributions),
                "elapsed": time.time() - self.current_round.started_at,
                "model_hash": self.current_round.model_hash,
            } if self.current_round else None,
            "recent_rounds": [
                {
                    "round_id": r.round_id,
                    "contributions": len(r.contributions),
                    "avg_loss": r.avg_loss,
                    "total_batch_size": r.total_batch_size,
                }
                for r in self.round_history[-10:]
            ]
        }


class GenesisDataLoader:
    """
    Loads training data from the verified Genesis Dataset.
    
    Features:
    - Dynamic shard count (reads from manifest)
    - User-configurable storage limit (max_storage_mb)
    - Shard rotation (cycles through dataset over time)
    - Multi-shard support (downloads multiple shards up to storage limit)
    - ASYNC PREFETCHING: Pre-downloads next shard while training on current
    
    Active only for nodes holding Layer 0 (Embedding Layer).
    
    Data Source: CloudFront CDN (backed by S3)
    """
    # CloudFront CDN URL - single source of truth (cached, DDoS protected)
    GENESIS_CDN_URL = "https://dwquwt9gkkeil.cloudfront.net"
    # Size per shard in MB (must match populate_genesis_s3.py)
    SHARD_SIZE_MB = 10
    
    def __init__(
        self, 
        node_id: str, 
        tokenizer, 
        cache_dir: str = None,  # Default to ~/.neuroshard/data_cache
        max_storage_mb: float = 100.0,  # User-configurable limit
        manifest_version: int = 1
    ):
        self.node_id = node_id
        self.tokenizer = tokenizer
        
        # Default cache_dir to ~/.neuroshard/data_cache for consistent storage
        if cache_dir is None:
            cache_dir = os.path.join(os.path.expanduser("~"), ".neuroshard", "data_cache")
        self.cache_dir = cache_dir
        self.max_storage_mb = max_storage_mb
        self.manifest_version = manifest_version
        
        # CloudFront CDN manifest URL - single source of truth
        self.manifest_url = f"{self.GENESIS_CDN_URL}/manifest.json"
        
        # Manifest data (cached, refreshed periodically)
        self.manifest = None
        self.total_shards = 0
        self.manifest_last_fetch = 0
        self.MANIFEST_REFRESH_INTERVAL = 3600  # Refresh manifest every hour
        
        # Shard management
        self.max_shards = max(1, int(max_storage_mb / self.SHARD_SIZE_MB))
        self.assigned_shard_ids = []  # List of shard IDs this node is responsible for
        self.loaded_shards = {}  # shard_id -> tensor data
        self.current_shard_idx = 0  # Index into assigned_shard_ids for rotation
        self.shard_rotation_count = 0  # How many times we've rotated through
        self.loading_shards = set()  # Track shards currently being downloaded
        self._shard_lock = threading.Lock()  # Lock for shard loading
        self._download_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="shard-download")
        
        # ASYNC PREFETCHING: Keep next shard ready in background
        self._prefetch_in_progress = set()  # Shard IDs being prefetched
        self._prefetch_ready = {}  # shard_id -> tensor data (ready to use)
        
        # Initialize Data Swarm for P2P downloading
        self.swarm = None 
        
        self.current_dataset = None
        self.dataset_iterator = 0
        
        # Fetch manifest and assign initial shards
        self._refresh_manifest()
        self._assign_shards()
        
        # Try to load learned tokenizer from CDN (for proper vocab)
        self._load_learned_tokenizer()
        
        # THUNDERING HERD PREVENTION: Add random jitter before first download
        # This spreads load across the CDN when many nodes start simultaneously
        # Jitter: 0-5 seconds based on node_id hash
        import random
        jitter_seed = int(hashlib.sha256(self.node_id.encode()).hexdigest()[:8], 16)
        jitter_seconds = (jitter_seed % 5000) / 1000.0  # 0-5 seconds
        
        def delayed_prefetch():
            time.sleep(jitter_seconds)
            self._start_prefetch_next()
        
        # Start prefetching first shard with jitter (non-blocking)
        threading.Thread(target=delayed_prefetch, daemon=True).start()
        
        logger.info(f"GenesisDataLoader initialized: "
                   f"total_shards={self.total_shards}, "
                   f"max_storage={max_storage_mb}MB ({self.max_shards} shards), "
                   f"assigned={self.assigned_shard_ids[:5]}{'...' if len(self.assigned_shard_ids) > 5 else ''}, "
                   f"prefetch_jitter={jitter_seconds:.2f}s")

    def _load_learned_tokenizer(self):
        """
        Download and load the learned tokenizer from Genesis CDN.
        Checks if the network has learned more tokens and updates locally.
        """
        try:
            tokenizer_url = f"{self.GENESIS_CDN_URL}/tokenizer.json"
            tokenizer_cache_path = os.path.join(self.cache_dir, "tokenizer.json")
            
            # Always try to fetch latest from CDN
            try:
                logger.debug(f"[GENESIS] Checking for tokenizer updates from {tokenizer_url}...")
                resp = requests.get(tokenizer_url, timeout=10)
                
                if resp.status_code == 200:
                    remote_tokenizer_data = resp.json()
                    remote_vocab_size = remote_tokenizer_data.get("next_merge_id", 0)
                    
                    # Check against current tokenizer
                    if remote_vocab_size > self.tokenizer.next_merge_id:
                        logger.info(f"[GENESIS] Found improved tokenizer! ({self.tokenizer.next_merge_id} -> {remote_vocab_size} tokens)")
                        
                        # Save to cache
                        os.makedirs(self.cache_dir, exist_ok=True)
                        with open(tokenizer_cache_path, 'w') as f:
                            f.write(resp.text)
                            
                        # Update our tokenizer's merges
                        from neuroshard.core.model.tokenizer import NeuroTokenizer
                        learned_tokenizer = NeuroTokenizer.load(tokenizer_cache_path)
                        
                        self.tokenizer.merges = learned_tokenizer.merges
                        self.tokenizer.merge_to_tokens = learned_tokenizer.merge_to_tokens
                        self.tokenizer.next_merge_id = learned_tokenizer.next_merge_id
                        
                        logger.info(f"[GENESIS] Tokenizer updated: {self.tokenizer.next_merge_id} tokens")
                        return
            except Exception as e:
                logger.debug(f"[GENESIS] Failed to check for tokenizer updates: {e}")
            
            # Fallback to cached version if download failed
            if os.path.exists(tokenizer_cache_path) and self.tokenizer.next_merge_id <= 266:
                logger.info(f"[GENESIS] Loading cached tokenizer from {tokenizer_cache_path}")
                try:
                    from neuroshard.core.model.tokenizer import NeuroTokenizer
                    learned_tokenizer = NeuroTokenizer.load(tokenizer_cache_path)
                    
                    if learned_tokenizer.next_merge_id > self.tokenizer.next_merge_id:
                        self.tokenizer.merges = learned_tokenizer.merges
                        self.tokenizer.merge_to_tokens = learned_tokenizer.merge_to_tokens
                        self.tokenizer.next_merge_id = learned_tokenizer.next_merge_id
                        logger.info(f"[GENESIS] Loaded cached tokenizer: {self.tokenizer.next_merge_id} tokens")
                except Exception as e:
                    logger.warning(f"[GENESIS] Failed to load cached tokenizer: {e}")

        except Exception as e:
            logger.warning(f"[GENESIS] Error managing tokenizer: {e}")

    def _refresh_manifest_sync(self):
        """Synchronous manifest fetch (runs in background thread)."""
        try:
            logger.info(f"[GENESIS] Fetching manifest from {self.manifest_url}...")
            resp = requests.get(self.manifest_url, timeout=15)
            if resp.status_code == 200:
                manifest_data = resp.json()
                total_shards = manifest_data.get("total_shards", 0)
                
                # Update state atomically
                with self._shard_lock:
                    self.manifest = manifest_data
                    self.total_shards = total_shards
                    self.manifest_last_fetch = time.time()
                
                logger.info(f"[GENESIS] Manifest loaded: {self.total_shards} shards available")
                
                # Also check if tokenizer has improved (in background)
                self._load_learned_tokenizer()
            else:
                logger.error(f"[GENESIS] Failed to fetch manifest: HTTP {resp.status_code}")
                logger.error(f"[GENESIS] Response: {resp.text[:200]}")
        except Exception as e:
            logger.error(f"[GENESIS] Failed to fetch manifest from {self.manifest_url}: {type(e).__name__}: {e}")
            import traceback
            logger.error(f"[GENESIS] Traceback: {traceback.format_exc()}")
    
    def _refresh_manifest(self):
        """Fetch latest manifest from S3 (non-blocking after first load)."""
        now = time.time()
        
        # First time initialization - must be synchronous
        if self.manifest is None:
            self._refresh_manifest_sync()
            if self.total_shards == 0:
                raise RuntimeError(f"Cannot fetch manifest from {self.manifest_url}. Check S3 bucket.")
            return
        
        # Subsequent refreshes - use cached if recent
        if (now - self.manifest_last_fetch) < self.MANIFEST_REFRESH_INTERVAL:
            return  # Use cached manifest
        
        # Refresh in background (non-blocking)
        self._download_executor.submit(self._refresh_manifest_sync)

    def _assign_shards(self):
        """
        Assign shards to this node based on:
        1. Node's deterministic hash (ensures different nodes get different shards)
        2. User's storage limit (max_shards)
        3. Rotation offset (allows cycling through entire dataset over time)
        """
        if self.total_shards == 0:
            self.assigned_shard_ids = [0]
            return
        
        # Base offset from node ID (deterministic)
        node_hash = int(hashlib.sha256(self.node_id.encode()).hexdigest(), 16)
        base_offset = node_hash % self.total_shards
        
        # Rotation offset (changes over time to cover more data)
        rotation_offset = (self.shard_rotation_count * self.max_shards) % self.total_shards
        
        # Assign shards starting from (base + rotation) offset
        self.assigned_shard_ids = []
        for i in range(self.max_shards):
            shard_id = (base_offset + rotation_offset + i) % self.total_shards
            self.assigned_shard_ids.append(shard_id)
        
        logger.info(f"Assigned {len(self.assigned_shard_ids)} shards: "
                   f"{self.assigned_shard_ids[:5]}{'...' if len(self.assigned_shard_ids) > 5 else ''}")

    def rotate_shards(self):
        """
        Rotate to next set of shards.
        Call this periodically to train on different parts of the dataset.
        """
        # Clear old loaded shards to free memory
        old_shards = list(self.loaded_shards.keys())
        self.loaded_shards.clear()
        self.current_dataset = None
        self.dataset_iterator = 0
        
        # Increment rotation counter
        self.shard_rotation_count += 1
        
        # Refresh manifest (in case new shards were added)
        self._refresh_manifest()
        
        # Reassign shards with new rotation offset
        self._assign_shards()
        
        # Clean up old shard files from disk
        self._cleanup_old_shards(old_shards)
        
        logger.info(f"Rotated to new shards (rotation #{self.shard_rotation_count})")

    def _cleanup_old_shards(self, old_shard_ids: list):
        """Remove old shard files from disk to stay within storage limit."""
        for shard_id in old_shard_ids:
            if shard_id not in self.assigned_shard_ids:
                shard_path = os.path.join(self.cache_dir, f"genesis_shard_{shard_id}.pt")
                try:
                    if os.path.exists(shard_path):
                        os.remove(shard_path)
                        logger.debug(f"Cleaned up old shard: {shard_path}")
                except Exception as e:
                    logger.warning(f"Failed to cleanup shard {shard_id}: {e}")

    def set_swarm(self, swarm):
        """Set the DataSwarm instance."""
        self.swarm = swarm
    
    def _start_prefetch_next(self):
        """Start prefetching the next shard(s) in background."""
        if not self.assigned_shard_ids:
            return
        
        # Prefetch current and next shard
        shards_to_prefetch = []
        for offset in [0, 1]:  # Current and next
            idx = (self.current_shard_idx + offset) % len(self.assigned_shard_ids)
            shard_id = self.assigned_shard_ids[idx]
            
            with self._shard_lock:
                # Skip if already loaded, prefetching, or ready
                if (shard_id in self.loaded_shards or 
                    shard_id in self._prefetch_in_progress or
                    shard_id in self._prefetch_ready or
                    shard_id in self.loading_shards):
                    continue
                
                shards_to_prefetch.append(shard_id)
                self._prefetch_in_progress.add(shard_id)
        
        # Start downloads in background
        for shard_id in shards_to_prefetch:
            target_url = self.get_shard_url(shard_id)
            logger.debug(f"Prefetching shard {shard_id} in background...")
            self._download_executor.submit(self._prefetch_shard_sync, shard_id, target_url)
    
    def _prefetch_shard_sync(self, shard_id: int, target_url: str):
        """Synchronous shard prefetch (runs in background thread)."""
        try:
            logger.info(f"[GENESIS] Downloading shard {shard_id}...")
            # Download the Shard
            shard_path = None
            
            if self.swarm:
                try:
                    shard_path = self.swarm.download_shard(shard_id, manifest_url=target_url)
                    logger.info(f"[GENESIS] Swarm download succeeded for shard {shard_id}")
                except Exception as e:
                    logger.warning(f"[GENESIS] Swarm prefetch failed: {e}")
            
            if not shard_path:
                logger.info(f"[GENESIS] Using HTTP fallback for shard {shard_id}")
                shard_path = self._http_fallback_download(shard_id, target_url)
                logger.info(f"[GENESIS] HTTP download completed for shard {shard_id}")
            
            # Load tensor into prefetch buffer
            tensor_data = torch.load(shard_path, weights_only=True)
            
            with self._shard_lock:
                # DYNAMIC MEMORY LIMIT: Based on user's max_storage_mb setting
                # Each shard is ~10MB compressed on disk, ~100-200MB uncompressed in RAM
                # Calculate max shards we can keep in memory
                shard_size_mb = 150  # Conservative estimate per shard in RAM
                max_cached_shards = max(3, int(self.max_storage_mb / shard_size_mb))
                
                total_loaded = len(self.loaded_shards) + len(self._prefetch_ready)
                if total_loaded >= max_cached_shards:
                    # Clear oldest loaded shard (not the current one)
                    current_shard = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)] if self.assigned_shard_ids else None
                    for old_shard_id in list(self.loaded_shards.keys()):
                        if old_shard_id != current_shard:
                            del self.loaded_shards[old_shard_id]
                            logger.debug(f"Evicted shard {old_shard_id} from cache (limit: {max_cached_shards} shards)")
                            break
                
                self._prefetch_ready[shard_id] = tensor_data
                self._prefetch_in_progress.discard(shard_id)
            
            logger.info(f"[GENESIS] Shard {shard_id} ready: {len(tensor_data):,} tokens")
            
        except Exception as e:
            logger.error(f"[GENESIS] Download FAILED for shard {shard_id}: {type(e).__name__}: {e}")
            import traceback
            logger.error(f"[GENESIS] Traceback: {traceback.format_exc()}")
            with self._shard_lock:
                self._prefetch_in_progress.discard(shard_id)
    
    def is_data_ready(self) -> bool:
        """Check if data is ready for training (non-blocking check)."""
        # Try to acquire lock with timeout to prevent blocking training loop
        acquired = self._shard_lock.acquire(timeout=0.5)
        if not acquired:
            # Lock held by download thread - assume data might be ready soon
            logger.debug("[GENESIS] Lock contention in is_data_ready - skipping check")
            return False
        
        try:
            # Data ready if we have current dataset OR prefetched shard is ready
            if self.current_dataset is not None and len(self.current_dataset) > 0:
                return True
            
            # Check if ANY assigned shard is ready (not just current)
            # This handles the case where prefetch completes before is_data_ready is called
            if self._prefetch_ready:
                # A prefetched shard is ready - we can use it
                return True
            
            # Also check loaded_shards
            if self.loaded_shards:
                return True
            
            # Check if current shard is specifically ready
            if self.assigned_shard_ids:
                shard_id = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)]
                if shard_id in self._prefetch_ready:
                    return True
                if shard_id in self.loaded_shards:
                    return True
            
            return False
        finally:
            self._shard_lock.release()

    def get_shard_url(self, shard_id: int) -> str:
        """Get download URL for a specific shard (always use CDN)."""
        # Always use CDN URL regardless of what manifest says
        # This ensures we go through CloudFront for caching/security
        return f"{self.GENESIS_CDN_URL}/shard_{shard_id}.pt"

    def _load_shard_sync(self, shard_id: int, target_url: str):
        """Synchronous shard loading (runs in background thread)."""
        # Download the Shard (Swarm or HTTP)
        shard_path = None
        
        if self.swarm:
            try:
                shard_path = self.swarm.download_shard(shard_id, manifest_url=target_url)
            except Exception as e:
                logger.error(f"Swarm download failed: {e}")
        
        if not shard_path:
            shard_path = self._http_fallback_download(shard_id, target_url)
                   
        # Load tensor
        try:
            tensor_data = torch.load(shard_path, weights_only=True)
            with self._shard_lock:
                self.loaded_shards[shard_id] = tensor_data
                self.current_dataset = tensor_data
                self.dataset_iterator = 0
                self.loading_shards.discard(shard_id)
            logger.info(f"Loaded Shard {shard_id}: {len(tensor_data)} tokens")
        except Exception as e:
            logger.error(f"Failed to load shard {shard_path}: {e}")
            with self._shard_lock:
                self.loading_shards.discard(shard_id)
                # Create dummy data if all else fails (use valid byte tokens 10-265)
                self.current_dataset = torch.randint(10, 266, (10000,), dtype=torch.long)

    def ensure_shard_loaded(self, shard_id: int = None):
        """
        Download and load a shard if not present.
        Opportunistically switches to ANY ready shard if the target isn't ready.
        """
        target_shard_id = shard_id
        
        if target_shard_id is None:
            # Default: try current shard in rotation
            if not self.assigned_shard_ids:
                return
            target_shard_id = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)]
        
        with self._shard_lock:
            # 1. Check if target is ready (Fastest)
            if target_shard_id in self.loaded_shards:
                self.current_dataset = self.loaded_shards[target_shard_id]
                return

            # 2. Check if target is in prefetch buffer
            if target_shard_id in self._prefetch_ready:
                self.current_dataset = self._prefetch_ready.pop(target_shard_id)
                self.loaded_shards[target_shard_id] = self.current_dataset
                self.dataset_iterator = 0
                logger.info(f"Using prefetched shard {target_shard_id}: {len(self.current_dataset)} tokens")
                self._start_prefetch_next_unlocked()
                return
            
            # 3. OPPORTUNISTIC: If target isn't ready, check if ANY assigned shard is ready in prefetch
            # This prevents blocking on shard A when shard B is already downloaded
            if shard_id is None:  # Only if caller didn't request specific shard
                for ready_id in list(self._prefetch_ready.keys()):
                    if ready_id in self.assigned_shard_ids:
                        # Switch to this ready shard!
                        logger.info(f"Opportunistically switching to ready shard {ready_id} (was waiting for {target_shard_id})")
                        
                        # Update index to match
                        try:
                            new_idx = self.assigned_shard_ids.index(ready_id)
                            self.current_shard_idx = new_idx
                        except ValueError:
                            pass
                            
                        self.current_dataset = self._prefetch_ready.pop(ready_id)
                        self.loaded_shards[ready_id] = self.current_dataset
                        self.dataset_iterator = 0
                        self._start_prefetch_next_unlocked()
                        return

            # 4. If still nothing, trigger download for target
            if target_shard_id in self.loading_shards or target_shard_id in self._prefetch_in_progress:
                logger.debug(f"Shard {target_shard_id} is already being downloaded, waiting...")
                return  # Don't block
            
            # Mark as loading and start download in background
            self.loading_shards.add(target_shard_id)
        
        target_url = self.get_shard_url(target_shard_id)
        logger.info(f"Loading Shard {target_shard_id} from {target_url}")
        
        # Submit to thread pool (non-blocking)
        self._download_executor.submit(self._load_shard_sync, target_shard_id, target_url)
    
    def _start_prefetch_next_unlocked(self):
        """Start prefetching next shard (call only when holding _shard_lock)."""
        # Schedule prefetch in background (don't hold lock during download)
        self._download_executor.submit(self._start_prefetch_next)

    def _http_fallback_download(self, shard_id: int, target_url: str = None) -> str:
        """Download shard from CloudFront CDN."""
        os.makedirs(self.cache_dir, exist_ok=True)
        shard_path = os.path.join(self.cache_dir, f"genesis_shard_{shard_id}.pt")
        
        if os.path.exists(shard_path):
            return shard_path
            
        # Use target URL from manifest, or construct CDN URL
        url = target_url or f"{self.GENESIS_CDN_URL}/shard_{shard_id}.pt"
        
        try:
            with requests.get(url, stream=True, timeout=60) as r:
                r.raise_for_status()
                with open(shard_path, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192): 
                        f.write(chunk)
            logger.info(f"Downloaded shard {shard_id}: {os.path.getsize(shard_path)/1e6:.1f}MB")
            return shard_path
        except Exception as e:
            logger.error(f"Failed to download shard {shard_id} from {url}: {e}")
            raise RuntimeError(f"Failed to download shard {shard_id}: {e}")

    def get_batch(self, batch_size: int = 4, seq_len: int = 512) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get a batch from the current shard.
        
        NON-BLOCKING VERSION: Returns quickly if data not ready.
        Uses prefetch buffer for instant shard switches.
        
        Automatically rotates to next shard when current one is exhausted.
        Returns (input_ids, labels).
        
        Raises RuntimeError if data not ready (caller should retry later).
        """
        # Try to load from prefetch buffer first
        self.ensure_shard_loaded()
        
        # NON-BLOCKING: Check if data is actually ready
        # Don't wait/block - let the caller handle the retry
        if self.current_dataset is None:
            # Check if anything is in progress
            with self._shard_lock:
                loading_any = bool(self.loading_shards or self._prefetch_in_progress)
                prefetch_ready = bool(self._prefetch_ready)
            
            if prefetch_ready:
                # There's a prefetched shard - try to use it
                self.ensure_shard_loaded()
            elif not loading_any:
                # Nothing loading - kick off a new load
                self._start_prefetch_next()
            
            # Return early - data not ready yet
            raise RuntimeError("Data not ready - shard still loading")
        
        data_len = len(self.current_dataset)
        req_len = (batch_size * seq_len) + 1 
        
        # Check if we've exhausted current shard
        if self.dataset_iterator + req_len > data_len:
            # Log completion of current shard
            completed_shard = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)]
            steps_done = data_len // req_len
            logger.info(f"✓ Completed shard {completed_shard} ({steps_done} steps, {data_len:,} tokens)")
            
            # Move to next shard in our assigned list
            self.current_shard_idx += 1
            
            if self.current_shard_idx >= len(self.assigned_shard_ids):
                # We've gone through all assigned shards - rotate to new set
                logger.info(f"Exhausted all {len(self.assigned_shard_ids)} assigned shards. Rotating to new set...")
                self.rotate_shards()
            
            # Try to use prefetched shard (FAST PATH)
            next_shard_id = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)]
            
            with self._shard_lock:
                if next_shard_id in self._prefetch_ready:
                    # Instant switch to prefetched shard
                    self.current_dataset = self._prefetch_ready.pop(next_shard_id)
                    self.loaded_shards[next_shard_id] = self.current_dataset
                    logger.info(f"Switched to prefetched shard {next_shard_id}: {len(self.current_dataset)} tokens")
                elif next_shard_id in self.loaded_shards:
                    self.current_dataset = self.loaded_shards[next_shard_id]
                else:
                    # Need to wait for next shard - trigger load
                    self.ensure_shard_loaded(next_shard_id)
                    raise RuntimeError("Data not ready - loading next shard")
            
            # Start prefetching the shard after next
            self._start_prefetch_next()
            
            self.dataset_iterator = 0
            data_len = len(self.current_dataset)
        
        start_idx = self.dataset_iterator
        end_idx = start_idx + req_len
        
        chunk = self.current_dataset[start_idx:end_idx]
        self.dataset_iterator += req_len
        
        # Log shard progress periodically (every 100 steps within shard)
        steps_in_shard = self.dataset_iterator // req_len
        total_steps_in_shard = data_len // req_len
        if steps_in_shard % 100 == 0:
            progress_pct = (self.dataset_iterator / data_len) * 100
            current_shard = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)]
            logger.info(f"Shard {current_shard} progress: {progress_pct:.1f}% "
                       f"({steps_in_shard}/{total_steps_in_shard} steps)")
        
        # Prepare batch
        exact_len = batch_size * seq_len
        
        inputs = chunk[:exact_len].view(batch_size, seq_len)
        labels = chunk[1:exact_len+1].view(batch_size, seq_len)
        
        return inputs, labels
    
    def get_stats(self) -> dict:
        """Get loader statistics."""
        # Calculate progress within current shard
        shard_progress = 0.0
        steps_in_shard = 0
        total_steps_in_shard = 0
        current_shard_id = None
        
        if self.current_dataset is not None and len(self.current_dataset) > 0:
            data_len = len(self.current_dataset)
            req_len = 1025  # Approximate: batch_size * seq_len + 1
            shard_progress = (self.dataset_iterator / data_len) * 100
            steps_in_shard = self.dataset_iterator // req_len
            total_steps_in_shard = data_len // req_len
            if self.assigned_shard_ids:
                current_shard_id = self.assigned_shard_ids[self.current_shard_idx % len(self.assigned_shard_ids)]
        
        return {
            "total_shards_available": self.total_shards,
            "max_shards_configured": self.max_shards,
            "max_storage_mb": self.max_storage_mb,
            "assigned_shards": len(self.assigned_shard_ids),
            "loaded_shards": len(self.loaded_shards),
            "prefetch_in_progress": len(self._prefetch_in_progress),
            "current_shard_idx": self.current_shard_idx,
            "current_shard_id": current_shard_id,
            "shard_progress_pct": round(shard_progress, 1),
            "steps_in_shard": steps_in_shard,
            "total_steps_in_shard": total_steps_in_shard,
            "rotation_count": self.shard_rotation_count,
            "storage_used_mb": len(self.loaded_shards) * self.SHARD_SIZE_MB,
        }


class DataValidator:
    """
    Validates training data quality before it enters the buffer.
    
    Prevents garbage/spam from polluting the local training set.
    """
    def __init__(self):
        pass
        
    def validate_text(self, text: str) -> Tuple[bool, str]:
        """
        Validate text quality.
        Returns (is_valid, reason).
        """
        if not text or not text.strip():
            return False, "Empty text"
            
        if len(text) < 20:
            return False, "Text too short (<20 chars)"
            
        # Entropy check (compression ratio)
        # Highly repetitive text compresses too well (ratio > 5.0)
        # Random text compresses poorly (ratio ~ 1.0)
        import zlib
        compressed = zlib.compress(text.encode())
        ratio = len(text) / len(compressed)
        
        if ratio > 6.0:
            return False, f"High compression ratio ({ratio:.1f}) - likely repetitive spam"
            
        if ratio < 1.1 and len(text) > 200:
            return False, f"Low compression ratio ({ratio:.1f}) - likely random gibberish"
            
        # Basic character distribution check
        # Check if text is mostly special characters
        alnum_count = sum(c.isalnum() for c in text)
        if alnum_count / len(text) < 0.5:
            return False, "Too many special characters"
            
        return True, "OK"


class FederatedDataManager:
    """
    Manages federated dataset for distributed training.
    
    Nodes can contribute:
    1. Text data (tokenized)
    2. Curated datasets
    3. Synthetic data from other models
    
    Privacy features:
    - Differential privacy (noise injection)
    - Data hashing (no raw text stored)
    - Local processing only
    """
    
    def __init__(self, tokenizer, max_seq_len: int = 2048):
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        
        # Validator
        self.validator = DataValidator()
        
        # Local data buffer
        self.data_buffer: List[torch.Tensor] = []
        self.max_buffer_size = 10000
        
        # Stats
        self.total_samples_contributed = 0
        self.total_tokens_contributed = 0
        self.rejected_samples = 0
    
    def add_text(self, text: str, apply_dp: bool = True, epsilon: float = 1.0):
        """
        Add text to the local training buffer.
        
        Args:
            text: Raw text to add
            apply_dp: Apply differential privacy
            epsilon: DP epsilon (lower = more private)
        """
        # Validate first
        is_valid, reason = self.validator.validate_text(text)
        if not is_valid:
            logger.warning(f"Rejected training data: {reason}")
            self.rejected_samples += 1
            return
            
        # Tokenize
        tokens = self.tokenizer.encode(text)
        
        if len(tokens) == 0:
            return
        
        # Chunk into sequences with overlap
        # Use smaller chunk size for flexibility
        chunk_size = min(self.max_seq_len, 512)  # Use 512 for training efficiency
        stride = chunk_size // 2  # 50% overlap
        
        chunks_added = 0
        for i in range(0, max(1, len(tokens) - chunk_size + 1), stride):
            chunk = tokens[i:i + chunk_size]
            
            # Pad if needed
            if len(chunk) < chunk_size:
                chunk = chunk + [self.tokenizer.pad_token_id] * (chunk_size - len(chunk))
            
            tensor = torch.tensor(chunk, dtype=torch.long)
            
            # Apply differential privacy (token-level noise)
            if apply_dp:
                tensor = self._apply_dp(tensor, epsilon)
            
            self.data_buffer.append(tensor)
            self.total_samples_contributed += 1
            self.total_tokens_contributed += len(chunk)
            chunks_added += 1
        
        # Also handle short texts (< chunk_size)
        if len(tokens) < chunk_size and chunks_added == 0:
            chunk = tokens + [self.tokenizer.pad_token_id] * (chunk_size - len(tokens))
            tensor = torch.tensor(chunk, dtype=torch.long)
            if apply_dp:
                tensor = self._apply_dp(tensor, epsilon)
            self.data_buffer.append(tensor)
            self.total_samples_contributed += 1
            self.total_tokens_contributed += len(tokens)
        
        # Trim buffer if too large
        if len(self.data_buffer) > self.max_buffer_size:
            self.data_buffer = self.data_buffer[-self.max_buffer_size:]
    
    def _apply_dp(self, tokens: torch.Tensor, epsilon: float) -> torch.Tensor:
        """Apply differential privacy to tokens."""
        # Simple DP: randomly replace some tokens
        # More sophisticated methods would use the exponential mechanism
        noise_mask = torch.rand(tokens.shape) < (1.0 / epsilon)
        # Use current_vocab_size (not max vocab_size) to only sample valid tokens
        valid_vocab_size = getattr(self.tokenizer, 'current_vocab_size', 266)
        random_tokens = torch.randint(0, valid_vocab_size, tokens.shape)
        return torch.where(noise_mask, random_tokens, tokens)
    
    def get_batch(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get a batch for training.
        
        Returns:
            (input_ids, labels) - labels are shifted input_ids
        """
        if len(self.data_buffer) < batch_size:
            raise ValueError(f"Not enough data: have {len(self.data_buffer)}, need {batch_size}")
        
        # Random sample
        import random
        indices = random.sample(range(len(self.data_buffer)), batch_size)
        batch = torch.stack([self.data_buffer[i] for i in indices])
        
        # For causal LM, labels = inputs shifted by 1
        input_ids = batch[:, :-1]
        labels = batch[:, 1:]
        
        return input_ids, labels
    
    def get_stats(self) -> Dict[str, Any]:
        """Get data contribution stats."""
        return {
            "buffer_size": len(self.data_buffer),
            "total_samples": self.total_samples_contributed,
            "total_tokens": self.total_tokens_contributed,
            "rejected_samples": self.rejected_samples,
        }

