//! JEPA Data Pipeline
//!
//! Complete data pipeline for I-JEPA / V-JEPA training:
//! - ImageRecord: decoded image storage
//! - MultiScaleRandomCrop: bilinear-resampled random crop augmentation
//! - RandomHorizontalFlip: stochastic horizontal flip
//! - JepaAugmentationPipeline: chained augmentation with ImageNet normalization
//! - JepaDataConfig: pipeline configuration
//! - InMemoryShard: in-memory image shard (CPU-testable without filesystem)
//! - JepaBatch: one ready-to-train batch with masks
//! - JepaDataPipeline: streaming batched iterator over shards
//! - DatasetStats: aggregate statistics
//! - WebDatasetShard: filesystem shard descriptor (stub)

// ============================================================================
// ImageRecord
// ============================================================================

/// A single decoded image stored in host memory.
#[derive(Debug, Clone)]
pub struct ImageRecord {
    /// Pixel data in HWC layout, values in [0, 1].
    pub pixels: Vec<f32>,
    pub width: usize,
    pub height: usize,
    pub channels: usize,
    /// Optional class label.
    pub label: Option<usize>,
    /// Dataset key / file name.
    pub key: String,
}

impl ImageRecord {
    /// Pixel value at (row, col, channel). Returns 0.0 if out of range.
    #[inline]
    fn get(&self, row: usize, col: usize, c: usize) -> f32 {
        if row < self.height && col < self.width && c < self.channels {
            self.pixels[row * self.width * self.channels + col * self.channels + c]
        } else {
            0.0
        }
    }
}

// ============================================================================
// LCG PRNG helper
// ============================================================================

#[inline(always)]
fn lcg_step(state: &mut u64) -> u64 {
    *state = state
        .wrapping_mul(6364136223846793005)
        .wrapping_add(1442695040888963407);
    *state
}

/// Sample a f32 in [0, 1) from the LCG state.
#[inline(always)]
fn lcg_f32(state: &mut u64) -> f32 {
    (lcg_step(state) >> 11) as f32 / (1u64 << 53) as f32
}

/// Sample a usize in [0, n).
#[inline(always)]
fn lcg_usize(state: &mut u64, n: usize) -> usize {
    if n == 0 {
        return 0;
    }
    (lcg_step(state) % n as u64) as usize
}

// ============================================================================
// MultiScaleRandomCrop
// ============================================================================

/// Random crop followed by bilinear resize to a fixed target size.
///
/// Mimics the torchvision `RandomResizedCrop` transform used in I-JEPA:
/// - Sample crop area as `scale * total_area` of the source image.
/// - Sample aspect ratio from `ratio_range`.
/// - Bilinear-resample to `(target_size, target_size, channels)`.
#[derive(Debug, Clone)]
pub struct MultiScaleRandomCrop {
    pub target_size: usize,
    pub scale_range: (f32, f32),
    pub ratio_range: (f32, f32),
}

impl MultiScaleRandomCrop {
    pub fn new(target_size: usize, scale_range: (f32, f32), ratio_range: (f32, f32)) -> Self {
        Self {
            target_size,
            scale_range,
            ratio_range,
        }
    }

    /// Crop and bilinear-resize a single `ImageRecord`.
    ///
    /// `seed` seeds the LCG PRNG so results are deterministic per sample.
    /// Returns a flat `[target_size, target_size, channels]` buffer.
    pub fn crop_and_resize(&self, image: &ImageRecord, seed: u64) -> Vec<f32> {
        let mut lcg = seed.wrapping_mul(2654435761).wrapping_add(1);

        let src_h = image.height;
        let src_w = image.width;
        let c = image.channels;
        let ts = self.target_size;

        // Clamp scale range to (0, 1]
        let scale_lo = self.scale_range.0.max(1e-6_f32).min(1.0);
        let scale_hi = self.scale_range.1.max(scale_lo).min(1.0);
        let ratio_lo = self.ratio_range.0.max(1e-4_f32);
        let ratio_hi = self.ratio_range.1.max(ratio_lo);

        // Try up to 10 attempts to find a valid crop; fall back to centre crop.
        let mut crop_top = 0usize;
        let mut crop_left = 0usize;
        let mut crop_h = src_h;
        let mut crop_w = src_w;

        for _ in 0..10 {
            let scale = scale_lo + lcg_f32(&mut lcg) * (scale_hi - scale_lo);
            let ratio = ratio_lo + lcg_f32(&mut lcg) * (ratio_hi - ratio_lo);

            let area = (src_h * src_w) as f32 * scale;
            let h_f = (area / ratio).sqrt();
            let w_f = (area * ratio).sqrt();
            let h = h_f.round() as usize;
            let w = w_f.round() as usize;

            if h == 0 || w == 0 || h > src_h || w > src_w {
                continue;
            }

            let top = lcg_usize(&mut lcg, src_h - h + 1);
            let left = lcg_usize(&mut lcg, src_w - w + 1);

            crop_top = top;
            crop_left = left;
            crop_h = h;
            crop_w = w;
            break;
        }

        // Bilinear resample crop → target_size × target_size
        let mut out = vec![0.0f32; ts * ts * c];
        for oy in 0..ts {
            for ox in 0..ts {
                // Map output pixel to source pixel in the crop region
                // Using half-pixel convention for better edge behaviour
                let src_y_f = (oy as f32 + 0.5) / ts as f32 * crop_h as f32 - 0.5 + crop_top as f32;
                let src_x_f =
                    (ox as f32 + 0.5) / ts as f32 * crop_w as f32 - 0.5 + crop_left as f32;

                let y0 = src_y_f.floor() as isize;
                let x0 = src_x_f.floor() as isize;
                let y1 = y0 + 1;
                let x1 = x0 + 1;

                let dy = src_y_f - y0 as f32;
                let dx = src_x_f - x0 as f32;

                // Clamp coordinates
                let y0c = y0.clamp(0, src_h as isize - 1) as usize;
                let y1c = y1.clamp(0, src_h as isize - 1) as usize;
                let x0c = x0.clamp(0, src_w as isize - 1) as usize;
                let x1c = x1.clamp(0, src_w as isize - 1) as usize;

                let base = (oy * ts + ox) * c;
                for ch in 0..c {
                    let p00 = image.get(y0c, x0c, ch);
                    let p01 = image.get(y0c, x1c, ch);
                    let p10 = image.get(y1c, x0c, ch);
                    let p11 = image.get(y1c, x1c, ch);
                    out[base + ch] = p00 * (1.0 - dy) * (1.0 - dx)
                        + p01 * (1.0 - dy) * dx
                        + p10 * dy * (1.0 - dx)
                        + p11 * dy * dx;
                }
            }
        }
        out
    }
}

// ============================================================================
// RandomHorizontalFlip
// ============================================================================

/// Randomly flip an image left-right with probability `prob`.
#[derive(Debug, Clone)]
pub struct RandomHorizontalFlip {
    pub prob: f32,
}

impl RandomHorizontalFlip {
    pub fn new(prob: f32) -> Self {
        Self { prob }
    }

    /// Apply (or not) horizontal flip.
    ///
    /// `pixels` must be in HWC layout with `w * channels` stride per row.
    pub fn apply(
        &self,
        pixels: &[f32],
        w: usize,
        h: usize,
        channels: usize,
        seed: u64,
    ) -> Vec<f32> {
        let mut lcg = seed
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        let r = lcg_f32(&mut lcg);

        if r >= self.prob {
            return pixels.to_vec();
        }

        let mut out = pixels.to_vec();
        for row in 0..h {
            for col in 0..w / 2 {
                let mirror = w - 1 - col;
                for c in 0..channels {
                    let a = row * w * channels + col * channels + c;
                    let b = row * w * channels + mirror * channels + c;
                    out.swap(a, b);
                }
            }
        }
        out
    }
}

// ============================================================================
// JepaAugmentationPipeline
// ============================================================================

/// Chained augmentation: MultiScaleRandomCrop → RandomHorizontalFlip → normalize.
#[derive(Debug, Clone)]
pub struct JepaAugmentationPipeline {
    pub crop: MultiScaleRandomCrop,
    pub flip: RandomHorizontalFlip,
    pub mean: [f32; 3],
    pub std: [f32; 3],
}

impl JepaAugmentationPipeline {
    /// Build from a `JepaDataConfig`.
    pub fn from_config(config: &JepaDataConfig) -> Self {
        let mean = if config.imagenet_normalize {
            [0.485, 0.456, 0.406]
        } else {
            [0.0, 0.0, 0.0]
        };
        let std = if config.imagenet_normalize {
            [0.229, 0.224, 0.225]
        } else {
            [1.0, 1.0, 1.0]
        };
        Self {
            crop: MultiScaleRandomCrop::new(
                config.image_size,
                config.scale_range,
                config.ratio_range,
            ),
            flip: RandomHorizontalFlip::new(if config.use_horizontal_flip { 0.5 } else { 0.0 }),
            mean,
            std,
        }
    }

    /// Per-channel ImageNet normalization in-place.
    pub fn normalize(pixels: &mut Vec<f32>, mean: [f32; 3], std: [f32; 3]) {
        let c = mean.len(); // 3
        for chunk in pixels.chunks_mut(c) {
            for (i, v) in chunk.iter_mut().enumerate() {
                let ch = i % c;
                let s = std[ch].max(1e-7_f32);
                *v = (*v - mean[ch]) / s;
            }
        }
    }

    /// Run the full augmentation pipeline on one `ImageRecord`.
    ///
    /// Returns a flat `[H_out * W_out * C]` buffer (C=channels of the source image).
    pub fn process(&self, image: &ImageRecord, seed: u64) -> Vec<f32> {
        let ts = self.crop.target_size;
        let c = image.channels;

        // 1. Multi-scale random crop → [ts, ts, c]
        let mut cropped = self.crop.crop_and_resize(image, seed);

        // 2. Horizontal flip (use a derived seed so crop and flip are independent)
        let flip_seed = seed
            .wrapping_mul(2654435761)
            .wrapping_add(1442695040888963407);
        let flipped = self.flip.apply(&cropped, ts, ts, c, flip_seed);
        cropped = flipped;

        // 3. Normalize per-channel
        Self::normalize(&mut cropped, self.mean, self.std);

        cropped
    }
}

// ============================================================================
// JepaDataConfig
// ============================================================================

/// Top-level configuration for the JEPA data pipeline.
#[derive(Debug, Clone)]
pub struct JepaDataConfig {
    /// Target image size after crop (e.g. 224).
    pub image_size: usize,
    /// Patch size used for masking (e.g. 16 → 14×14 grid on 224px images).
    pub patch_size: usize,
    /// Number of images per batch.
    pub batch_size: usize,
    /// Concurrency hint (not used in CPU reference implementation).
    pub num_workers: usize,
    /// Paths to WebDataset shards (.tar files).
    pub shard_paths: Vec<String>,
    /// Scale range for `MultiScaleRandomCrop` (area fraction).
    pub scale_range: (f32, f32),
    /// Aspect ratio range for `MultiScaleRandomCrop`.
    pub ratio_range: (f32, f32),
    /// Whether to apply random horizontal flip.
    pub use_horizontal_flip: bool,
    /// Whether to apply ImageNet mean/std normalization.
    pub imagenet_normalize: bool,
}

impl Default for JepaDataConfig {
    fn default() -> Self {
        Self {
            image_size: 224,
            patch_size: 16,
            batch_size: 1024,
            num_workers: 8,
            shard_paths: Vec::new(),
            scale_range: (0.2, 1.0),
            ratio_range: (0.75, 1.33),
            use_horizontal_flip: true,
            imagenet_normalize: true,
        }
    }
}

impl JepaDataConfig {
    /// Grid side length in patches.
    pub fn grid_size(&self) -> usize {
        self.image_size / self.patch_size
    }

    /// Total patch count per image.
    pub fn num_patches(&self) -> usize {
        let g = self.grid_size();
        g * g
    }
}

// ============================================================================
// InMemoryShard
// ============================================================================

/// A shard of images held entirely in host memory.
///
/// Used for unit testing without a real filesystem.
#[derive(Debug, Clone)]
pub struct InMemoryShard {
    pub records: Vec<ImageRecord>,
    pub shard_id: usize,
}

impl InMemoryShard {
    /// Create `n` synthetic images of size `image_size × image_size × 3`
    /// with LCG-generated random pixels in `[0, 1]` and sequential labels.
    pub fn synthetic(n: usize, image_size: usize, shard_id: usize) -> Self {
        let mut records = Vec::with_capacity(n);
        let mut lcg: u64 = (shard_id as u64)
            .wrapping_mul(2654435761)
            .wrapping_add(1442695040888963407);

        for i in 0..n {
            let pixel_count = image_size * image_size * 3;
            let pixels: Vec<f32> = (0..pixel_count).map(|_| lcg_f32(&mut lcg)).collect();

            records.push(ImageRecord {
                pixels,
                width: image_size,
                height: image_size,
                channels: 3,
                label: Some(i % 1000),
                key: format!("shard{shard_id}/img{i:06}"),
            });
        }

        Self { records, shard_id }
    }

    /// Number of records in this shard.
    pub fn len(&self) -> usize {
        self.records.len()
    }

    /// True if the shard is empty.
    pub fn is_empty(&self) -> bool {
        self.records.is_empty()
    }
}

// ============================================================================
// JepaBatch
// ============================================================================

/// One training batch ready for JEPA forward + loss.
#[derive(Debug, Clone)]
pub struct JepaBatch {
    /// Augmented image pixels for each sample, flattened to `[H*W*C]`.
    pub images: Vec<Vec<f32>>,
    /// Context patch indices per sample `[batch_size][n_context]`.
    pub context_patch_indices: Vec<Vec<usize>>,
    /// Target block patch indices per sample `[batch_size][num_blocks][n_patches]`.
    pub target_patch_groups: Vec<Vec<Vec<usize>>>,
    /// Union of all target patch indices per sample `[batch_size][n_total_target]`.
    pub all_target_indices: Vec<Vec<usize>>,
    /// Class labels (if available) per sample.
    pub labels: Vec<Option<usize>>,
    /// Monotonically increasing batch ID within the epoch.
    pub batch_id: usize,
}

impl JepaBatch {
    /// Actual number of samples in this batch (may be less than config.batch_size at epoch end).
    pub fn batch_size(&self) -> usize {
        self.images.len()
    }
}

// ============================================================================
// JepaDataPipeline
// ============================================================================

/// Streaming data pipeline: iterates over shards, applies augmentation, and
/// assembles `JepaBatch` instances with JEPA block masks.
pub struct JepaDataPipeline {
    pub config: JepaDataConfig,
    augmentation: JepaAugmentationPipeline,
    mask_strategy: super::jepa::BlockMaskStrategy,
    shards: Vec<InMemoryShard>,
    /// Order in which shards are visited this epoch (shuffled each epoch).
    shard_order: Vec<usize>,
    /// Index into `shard_order` for the currently active shard.
    current_shard: usize,
    /// Index within the current shard's records.
    current_idx: usize,
    /// Current epoch (starts at 0, increments on `reset_epoch`).
    epoch: usize,
    /// Monotonic batch counter (resets on `reset_epoch`).
    batch_counter: usize,
    /// LCG seed for shard shuffling.
    shuffle_lcg: u64,
}

impl JepaDataPipeline {
    /// Build a new pipeline.
    pub fn new(config: JepaDataConfig, shards: Vec<InMemoryShard>) -> Self {
        let augmentation = JepaAugmentationPipeline::from_config(&config);
        let mask_strategy = super::jepa::BlockMaskStrategy::default_ijepa();
        let n = shards.len();
        let shard_order: Vec<usize> = (0..n).collect();
        Self {
            config,
            augmentation,
            mask_strategy,
            shards,
            shard_order,
            current_shard: 0,
            current_idx: 0,
            epoch: 0,
            batch_counter: 0,
            shuffle_lcg: 0xdeadbeef_cafebabe_u64,
        }
    }

    /// Total images across all shards.
    fn total_images(&self) -> usize {
        self.shards.iter().map(|s| s.len()).sum()
    }

    /// Number of full (and potentially partial) batches per epoch.
    pub fn num_batches_per_epoch(&self) -> usize {
        let total = self.total_images();
        if total == 0 || self.config.batch_size == 0 {
            return 0;
        }
        (total + self.config.batch_size - 1) / self.config.batch_size
    }

    /// Current epoch number.
    pub fn epoch(&self) -> usize {
        self.epoch
    }

    /// Advance the shard pointer, skipping empty shards.
    /// Returns false if all shards for this epoch have been consumed.
    fn advance_shard(&mut self) -> bool {
        self.current_shard += 1;
        self.current_idx = 0;
        self.current_shard < self.shard_order.len()
    }

    /// Reference to the currently active shard, or None if epoch is done.
    fn active_shard(&self) -> Option<&InMemoryShard> {
        let si = self.shard_order.get(self.current_shard)?;
        self.shards.get(*si)
    }

    /// Pull the next `ImageRecord` from the stream, advancing shards as needed.
    fn next_record(&mut self) -> Option<&ImageRecord> {
        loop {
            let shard = match self.active_shard() {
                Some(s) => s,
                None => return None,
            };

            if self.current_idx < shard.len() {
                // This looks like a borrow issue but we borrow immutably here.
                // We will return a raw pointer trick via direct indexing below.
                break;
            } else {
                if !self.advance_shard() {
                    return None;
                }
            }
        }

        // Re-resolve to get the record (needed because we can't hold shard ref through advance_shard).
        let si = *self.shard_order.get(self.current_shard)?;
        let shard = self.shards.get(si)?;
        let rec = shard.records.get(self.current_idx)?;
        self.current_idx += 1;
        Some(rec)
    }

    /// Yield the next batch, or `None` when the epoch is exhausted.
    pub fn next_batch(&mut self) -> Option<JepaBatch> {
        let batch_size = self.config.batch_size;
        let grid = self.config.grid_size();
        let batch_id = self.batch_counter;

        // Collect up to `batch_size` records.
        // We need to collect references first, then process.
        // Because next_record borrows &mut self, we collect indices instead.
        let mut sample_indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); // (shard_order_idx, record_idx)

        for _ in 0..batch_size {
            loop {
                let shard = match self.active_shard() {
                    Some(s) => s,
                    None => break,
                };
                if self.current_idx < shard.len() {
                    let si = self.shard_order[self.current_shard];
                    sample_indices.push((si, self.current_idx));
                    self.current_idx += 1;
                    break;
                } else {
                    if !self.advance_shard() {
                        break;
                    }
                }
            }

            // Check if we are done.
            if self.active_shard().is_none()
                && sample_indices.len() < batch_size
                && !sample_indices.is_empty()
            {
                // Last partial batch — collect what we have.
                break;
            }

            if sample_indices.len() == batch_size {
                break;
            }
        }

        if sample_indices.is_empty() {
            return None;
        }

        let actual_bs = sample_indices.len();
        let mut images = Vec::with_capacity(actual_bs);
        let mut context_patch_indices = Vec::with_capacity(actual_bs);
        let mut target_patch_groups = Vec::with_capacity(actual_bs);
        let mut all_target_indices_out = Vec::with_capacity(actual_bs);
        let mut labels = Vec::with_capacity(actual_bs);

        for (sample_i, (shard_idx, record_idx)) in sample_indices.iter().enumerate() {
            let record = &self.shards[*shard_idx].records[*record_idx];

            // Build a per-sample seed from batch_id, sample_i, and shard/record position.
            let seed: u64 = (batch_id as u64)
                .wrapping_mul(1000003)
                .wrapping_add(*shard_idx as u64)
                .wrapping_mul(999983)
                .wrapping_add(*record_idx as u64)
                .wrapping_add(sample_i as u64);

            // Augment image.
            let aug = self.augmentation.process(record, seed);
            images.push(aug);
            labels.push(record.label);

            // Generate JEPA block masks.
            let mask = self.mask_strategy.generate(grid, grid, seed);
            context_patch_indices.push(mask.context_indices);
            target_patch_groups.push(mask.target_blocks);
            all_target_indices_out.push(mask.all_target_indices);
        }

        self.batch_counter += 1;

        Some(JepaBatch {
            images,
            context_patch_indices,
            target_patch_groups,
            all_target_indices: all_target_indices_out,
            labels,
            batch_id,
        })
    }

    /// Reset for the next epoch: shuffle shard order and rewind position counters.
    pub fn reset_epoch(&mut self) {
        self.epoch += 1;
        self.batch_counter = 0;
        self.current_shard = 0;
        self.current_idx = 0;

        // Fisher-Yates shuffle of shard order using internal LCG.
        let n = self.shard_order.len();
        for i in 0..n {
            self.shard_order[i] = i;
        }
        for i in 0..n {
            let j = i + lcg_usize(&mut self.shuffle_lcg, n - i);
            self.shard_order.swap(i, j);
        }
    }
}

// ============================================================================
// DatasetStats
// ============================================================================

/// Aggregate statistics about the dataset and masking behaviour.
#[derive(Debug, Clone)]
pub struct DatasetStats {
    pub total_images: usize,
    pub num_shards: usize,
    pub images_per_shard: Vec<usize>,
    /// Mean context patch count across the first sampled batch.
    pub avg_context_patches: f64,
    /// Mean target patch count across the first sampled batch.
    pub avg_target_patches: f64,
    /// Fraction of total patches that are target patches (mask_efficiency).
    pub mask_efficiency: f64,
}

impl DatasetStats {
    /// Compute statistics by sampling one batch from a separate mask strategy instance.
    pub fn compute(pipeline: &JepaDataPipeline) -> Self {
        let total_images: usize = pipeline.shards.iter().map(|s| s.len()).sum();
        let num_shards = pipeline.shards.len();
        let images_per_shard: Vec<usize> = pipeline.shards.iter().map(|s| s.len()).collect();

        let total_patches = pipeline.config.num_patches();
        let grid = pipeline.config.grid_size();

        // Sample a small number of masks to estimate average patch counts.
        let sample_size = 32.min(total_images.max(1));
        let mut mask_strategy = super::jepa::BlockMaskStrategy::default_ijepa();
        let masks = mask_strategy.generate_batch(sample_size, grid, grid);

        let avg_context = if masks.is_empty() {
            0.0
        } else {
            masks
                .iter()
                .map(|m| m.context_indices.len() as f64)
                .sum::<f64>()
                / masks.len() as f64
        };

        let avg_target = if masks.is_empty() {
            0.0
        } else {
            masks
                .iter()
                .map(|m| m.all_target_indices.len() as f64)
                .sum::<f64>()
                / masks.len() as f64
        };

        let mask_efficiency = if total_patches == 0 {
            0.0
        } else {
            avg_target / total_patches as f64
        };

        Self {
            total_images,
            num_shards,
            images_per_shard,
            avg_context_patches: avg_context,
            avg_target_patches: avg_target,
            mask_efficiency,
        }
    }
}

// ============================================================================
// WebDatasetShard
// ============================================================================

/// Filesystem shard descriptor for WebDataset-format `.tar` archives.
///
/// Use [`WebDatasetShard::load`] to read the actual archive from disk
/// (via [`read_webdataset_shard`]); [`WebDatasetShard::to_in_memory`]
/// generates synthetic data and exists for tests that need a shard
/// without touching the filesystem.
#[derive(Debug, Clone)]
pub struct WebDatasetShard {
    pub path: String,
    pub num_records: usize,
    pub shard_id: usize,
    pub compressed: bool,
}

impl WebDatasetShard {
    /// Create a descriptor pointing at `path`.
    ///
    /// `num_records` is used both for progress estimation and for synthetic
    /// data generation in `to_in_memory`.
    pub fn new(path: &str, num_records: usize, shard_id: usize) -> Self {
        Self {
            path: path.to_owned(),
            num_records,
            shard_id,
            compressed: path.ends_with(".tar.gz") || path.ends_with(".tgz"),
        }
    }

    /// Read the shard's `.tar` (or gzip-compressed `.tar.gz`/`.tgz`)
    /// archive from disk and decode its records into an [`InMemoryShard`].
    pub fn load(&self) -> Result<InMemoryShard, String> {
        let (raw_records, _stats) = read_webdataset_shard(std::path::Path::new(&self.path))?;
        let records: Vec<ImageRecord> = raw_records
            .into_iter()
            .map(webdataset_record_to_image)
            .collect();
        Ok(InMemoryShard {
            records,
            shard_id: self.shard_id,
        })
    }

    /// Generate a synthetic `InMemoryShard` with `num_records` records of
    /// size `image_size × image_size × 3`. Used for testing without a
    /// real filesystem — use [`WebDatasetShard::load`] for real data.
    pub fn to_in_memory(&self, image_size: usize) -> InMemoryShard {
        InMemoryShard::synthetic(self.num_records, image_size, self.shard_id)
    }
}

// ============================================================================
// WebDatasetRecord
// ============================================================================

/// A record decoded from a WebDataset shard (one image + optional label).
#[derive(Debug, Clone)]
pub struct WebDatasetRecord {
    /// Numeric prefix, e.g. "000042".
    pub key: String,
    /// Raw JPEG/PNG bytes (not decoded to pixels).
    pub image_bytes: Vec<u8>,
    /// Parsed from .cls file.
    pub label: Option<usize>,
    /// "jpg" or "png".
    pub extension: String,
}

// ============================================================================
// ShardLoadStats
// ============================================================================

/// Statistics for a loaded shard.
#[derive(Debug, Clone)]
pub struct ShardLoadStats {
    pub shard_id: usize,
    pub records_loaded: usize,
    pub bytes_read: u64,
    pub load_duration_ms: u64,
}

// ============================================================================
// ShuffleBuffer
// ============================================================================

/// Fixed-size ring buffer for dataset shuffling.
/// Holds up to `capacity` records; reservoir sampling on insert.
pub struct ShuffleBuffer {
    capacity: usize,
    buffer: Vec<ImageRecord>,
    rng_state: u64,
}

impl ShuffleBuffer {
    /// Create a new shuffle buffer seeded from `capacity`.
    pub fn new(capacity: usize) -> Self {
        let rng_state = (capacity as u64).wrapping_mul(6364136223846793005);
        Self {
            capacity,
            buffer: Vec::new(),
            rng_state,
        }
    }

    /// Insert a record. If the buffer is not full, push it; otherwise replace
    /// a random slot using the internal LCG.
    pub fn push(&mut self, record: ImageRecord) {
        if self.buffer.len() < self.capacity {
            self.buffer.push(record);
        } else if self.capacity > 0 {
            let slot = lcg_usize(&mut self.rng_state, self.capacity);
            self.buffer[slot] = record;
        }
    }

    /// Remove and return the first `n` items (or all items if fewer than `n` remain).
    pub fn drain_batch(&mut self, n: usize) -> Vec<ImageRecord> {
        let take = n.min(self.buffer.len());
        self.buffer.drain(..take).collect()
    }

    /// Number of records currently in the buffer.
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    /// True when `len() >= capacity`.
    pub fn is_full(&self) -> bool {
        self.len() >= self.capacity
    }
}

// ============================================================================
// Tar parsing helpers (stdlib only)
// ============================================================================

/// Parse a null-terminated octal ASCII string from a fixed-width tar field.
fn parse_octal(field: &[u8]) -> u64 {
    // Trim leading/trailing null bytes and spaces
    let s = field
        .iter()
        .take_while(|&&b| b != 0 && b != b' ')
        .copied()
        .collect::<Vec<u8>>();
    let s = std::str::from_utf8(&s).unwrap_or("0");
    u64::from_str_radix(s.trim(), 8).unwrap_or(0)
}

/// Parse all WebDataset records from an in-memory tar byte slice.
///
/// A tar archive is a sequence of 512-byte blocks:
/// - Header block: filename (bytes 0..100), size (bytes 124..136, octal),
///   typeflag (byte 156, '0' or '\0' = regular file, '5' = directory).
/// - Data blocks: ceil(size / 512) × 512 bytes follow the header.
/// - End-of-archive: two consecutive all-zero blocks.
pub fn parse_tar_bytes(data: &[u8]) -> Vec<WebDatasetRecord> {
    // Accumulate raw file entries keyed by stem (numeric prefix).
    // Values: (image_bytes, image_ext, label_bytes)
    let mut image_map: std::collections::HashMap<String, (Vec<u8>, String)> =
        std::collections::HashMap::new();
    let mut label_map: std::collections::HashMap<String, Vec<u8>> =
        std::collections::HashMap::new();

    let mut offset = 0usize;

    loop {
        // Need at least one header block.
        if offset + 512 > data.len() {
            break;
        }

        let header = &data[offset..offset + 512];
        offset += 512;

        // End-of-archive: two consecutive zero blocks. A zero block starts here.
        if header.iter().all(|&b| b == 0) {
            break;
        }

        // Parse filename (bytes 0..100, null-terminated).
        let fname_raw = &header[0..100];
        let fname_len = fname_raw.iter().position(|&b| b == 0).unwrap_or(100);
        let filename = match std::str::from_utf8(&fname_raw[..fname_len]) {
            Ok(s) => s.trim_matches('/').to_owned(),
            Err(_) => {
                // Skip unreadable header.
                continue;
            }
        };

        // Parse typeflag (byte 156). '5' = directory, skip data blocks but
        // don't advance data pointer here — size may still be non-zero.
        let typeflag = header[156];
        let is_directory = typeflag == b'5';

        // Parse size (bytes 124..136, octal ASCII).
        let size = parse_octal(&header[124..136]) as usize;

        // Skip aligned data blocks for directories and zero-size files.
        let data_blocks = (size + 511) / 512;
        let data_bytes = data_blocks * 512;

        if is_directory || size == 0 {
            offset += data_bytes;
            continue;
        }

        // Read `size` bytes of content.
        if offset + data_bytes > data.len() {
            // Truncated archive — stop.
            break;
        }
        let content = data[offset..offset + size].to_vec();
        offset += data_bytes;

        // Derive extension and key from the filename.
        // Filenames may include a directory prefix; take the basename.
        let basename = filename.rsplit('/').next().unwrap_or(&filename);

        let (stem, ext) = if let Some(dot) = basename.rfind('.') {
            (&basename[..dot], &basename[dot + 1..])
        } else {
            (basename, "")
        };

        match ext {
            "jpg" | "jpeg" | "png" => {
                let image_ext = if ext == "jpeg" { "jpg" } else { ext };
                image_map.insert(stem.to_owned(), (content, image_ext.to_owned()));
            }
            "cls" => {
                label_map.insert(stem.to_owned(), content);
            }
            _ => {
                // Unknown extension — ignore.
            }
        }
    }

    // Merge image + label by key.
    let mut keys: Vec<String> = image_map.keys().cloned().collect();
    keys.sort();

    keys.into_iter()
        .filter_map(|key| {
            let (image_bytes, extension) = image_map.remove(&key)?;
            let label = label_map.get(&key).and_then(|bytes| {
                let s = std::str::from_utf8(bytes).ok()?.trim().to_owned();
                s.parse::<usize>().ok()
            });
            Some(WebDatasetRecord {
                key,
                image_bytes,
                label,
                extension,
            })
        })
        .collect()
}

/// Read all records from a WebDataset tar shard at `path`.
///
/// Returns records in shard order plus load statistics.
/// Gzip-compressed shards (`.tar.gz`/`.tgz`, detected by the 1F 8B magic
/// bytes rather than extension) are decompressed transparently.
pub fn read_webdataset_shard(
    path: &std::path::Path,
) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
    use std::io::Read;
    use std::time::Instant;

    let t0 = Instant::now();

    let mut file =
        std::fs::File::open(path).map_err(|e| format!("cannot open {}: {e}", path.display()))?;

    let mut data = Vec::new();
    file.read_to_end(&mut data)
        .map_err(|e| format!("read error for {}: {e}", path.display()))?;

    let bytes_read = data.len() as u64;

    // Gzip magic: 0x1F 0x8B. Decompress before tar parsing.
    if data.len() >= 2 && data[0] == 0x1f && data[1] == 0x8b {
        let mut decoder = flate2::read::GzDecoder::new(&data[..]);
        let mut decompressed = Vec::new();
        decoder
            .read_to_end(&mut decompressed)
            .map_err(|e| format!("gzip decompression failed for {}: {e}", path.display()))?;
        data = decompressed;
    }

    let records = parse_tar_bytes(&data);
    let records_loaded = records.len();
    let load_duration_ms = t0.elapsed().as_millis() as u64;

    // shard_id is unknown at this level; callers should override if needed.
    let stats = ShardLoadStats {
        shard_id: 0,
        records_loaded,
        bytes_read,
        load_duration_ms,
    };

    Ok((records, stats))
}

/// Decode a `WebDatasetRecord` into an `ImageRecord`.
///
/// Pixel values are placeholder 0.5f32 until rtx-vision image decoder is
/// wired. The function detects JPEG (`\xFF\xD8`) and PNG (`\x89PNG`) magic
/// bytes to validate the format; other byte sequences also receive the
/// placeholder.
fn webdataset_record_to_image(rec: WebDatasetRecord) -> ImageRecord {
    const TARGET_W: usize = 224;
    const TARGET_H: usize = 224;
    const C: usize = 3;

    #[cfg(feature = "image-decode")]
    {
        use image::io::Reader as ImageReader;
        use std::io::Cursor;

        if let Ok(reader) = ImageReader::new(Cursor::new(&rec.image_bytes)).with_guessed_format() {
            if let Ok(img) = reader.decode() {
                // Resize to 224×224 and convert to RGB
                let rgb = img
                    .resize_exact(
                        TARGET_W as u32,
                        TARGET_H as u32,
                        image::imageops::FilterType::Triangle,
                    )
                    .into_rgb8();

                let pixels: Vec<f32> = rgb
                    .into_raw()
                    .into_iter()
                    .map(|p| p as f32 / 255.0)
                    .collect();

                return ImageRecord {
                    pixels,
                    width: TARGET_W,
                    height: TARGET_H,
                    channels: C,
                    label: rec.label,
                    key: rec.key,
                };
            }
        }
    }

    // Fallback: placeholder (without image-decode feature, or on decode failure)
    let _ = &rec.image_bytes; // suppress unused warning
    ImageRecord {
        pixels: vec![0.5f32; TARGET_W * TARGET_H * C],
        width: TARGET_W,
        height: TARGET_H,
        channels: C,
        label: rec.label,
        key: rec.key,
    }
}

impl JepaDataPipeline {
    /// Load a pipeline from filesystem shards.
    ///
    /// `shard_paths` is a list of `.tar` files (uncompressed WebDataset format).
    /// Each file is read, parsed, and converted to an `InMemoryShard`.
    ///
    /// Pixel values are placeholder 0.5f32 until rtx-vision image decoder is
    /// wired.
    pub fn from_filesystem(
        config: JepaDataConfig,
        shard_paths: Vec<std::path::PathBuf>,
    ) -> Result<Self, String> {
        // Validate all paths exist before loading anything.
        for p in &shard_paths {
            if !p.exists() {
                return Err(format!("shard path does not exist: {}", p.display()));
            }
        }

        let mut shards: Vec<InMemoryShard> = Vec::with_capacity(shard_paths.len());

        for (shard_id, path) in shard_paths.iter().enumerate() {
            let (raw_records, _stats) = read_webdataset_shard(path)?;

            let records: Vec<ImageRecord> = raw_records
                .into_iter()
                .map(webdataset_record_to_image)
                .collect();

            shards.push(InMemoryShard { records, shard_id });
        }

        Ok(Self::new(config, shards))
    }
}

// ============================================================================
// HTTP / WebDataset URL loading
// ============================================================================

/// Descriptor for a remote WebDataset shard at a URL.
#[derive(Debug, Clone)]
pub struct UrlShardDescriptor {
    /// Full HTTP/HTTPS URL to the .tar shard file
    pub url: String,
    /// Expected number of records (0 = unknown)
    pub expected_records: usize,
    /// Shard index in the dataset
    pub shard_id: usize,
}

impl UrlShardDescriptor {
    pub fn new(url: impl Into<String>, expected_records: usize, shard_id: usize) -> Self {
        Self {
            url: url.into(),
            expected_records,
            shard_id,
        }
    }

    /// Returns true if this descriptor points to an HTTPS URL (vs plain HTTP)
    pub fn is_https(&self) -> bool {
        self.url.starts_with("https://")
    }
}

/// Like `download_webdataset_shard` but with a configurable timeout.
/// Used internally and for testing.
pub(crate) fn download_webdataset_shard_with_timeout(
    url: &str,
    shard_id: usize,
    timeout_secs: u64,
) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
    let start = std::time::Instant::now();

    // Validate URL scheme
    if !url.starts_with("http://") && !url.starts_with("https://") {
        return Err(format!(
            "unsupported URL scheme (expected http:// or https://): {url}"
        ));
    }

    // Download using a one-shot Tokio runtime
    let bytes = {
        let rt = tokio::runtime::Runtime::new()
            .map_err(|e| format!("failed to create Tokio runtime: {e}"))?;

        rt.block_on(async {
            let client = reqwest::Client::builder()
                .timeout(std::time::Duration::from_secs(timeout_secs))
                .build()
                .map_err(|e| format!("reqwest client build failed: {e}"))?;

            let response = client
                .get(url)
                .send()
                .await
                .map_err(|e| format!("HTTP GET failed for {url}: {e}"))?;

            if !response.status().is_success() {
                return Err(format!("HTTP {} for {url}", response.status()));
            }

            response
                .bytes()
                .await
                .map(|b| b.to_vec())
                .map_err(|e| format!("failed to read response body: {e}"))
        })?
    };

    let load_duration_ms = start.elapsed().as_millis() as u64;
    let bytes_read = bytes.len() as u64;

    let records = parse_tar_bytes(&bytes);
    let records_loaded = records.len();

    let stats = ShardLoadStats {
        shard_id,
        records_loaded,
        bytes_read,
        load_duration_ms,
    };

    Ok((records, stats))
}

/// Download a WebDataset `.tar` shard from an HTTP/HTTPS URL and parse it.
///
/// Returns `(records, stats)` on success, `Err(message)` on failure.
/// Requires `reqwest` (already a workspace dep).
///
/// The download is synchronous (blocks via `tokio::runtime::Handle`). For
/// async callers, use the async variant `download_shard_async` instead.
pub fn download_webdataset_shard(
    url: &str,
    shard_id: usize,
) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
    download_webdataset_shard_with_timeout(url, shard_id, 300)
}

impl JepaDataPipeline {
    /// Load a pipeline from HTTP/HTTPS WebDataset shard URLs.
    ///
    /// Each URL must point to an uncompressed `.tar` file in WebDataset format.
    /// Downloads are sequential (one shard at a time).
    ///
    /// Returns `Err` if any URL fails to download or parse.
    pub fn from_urls(config: JepaDataConfig, urls: Vec<String>) -> Result<Self, String> {
        // Validate all URLs before downloading
        for url in &urls {
            if !url.starts_with("http://") && !url.starts_with("https://") {
                return Err(format!(
                    "unsupported URL scheme in '{}' (expected http:// or https://)",
                    url
                ));
            }
        }

        let mut shards: Vec<InMemoryShard> = Vec::with_capacity(urls.len());

        for (shard_id, url) in urls.iter().enumerate() {
            let (raw_records, _stats) = download_webdataset_shard(url, shard_id)?;

            let records: Vec<ImageRecord> = raw_records
                .into_iter()
                .map(webdataset_record_to_image)
                .collect();

            shards.push(InMemoryShard { records, shard_id });
        }

        Ok(Self::new(config, shards))
    }

    /// Validate a list of URLs without downloading.
    ///
    /// Returns `Ok(())` if all URLs have valid schemes, `Err` with the first invalid URL.
    pub fn validate_urls(urls: &[String]) -> Result<(), String> {
        for url in urls {
            if !url.starts_with("http://") && !url.starts_with("https://") {
                return Err(format!(
                    "invalid URL '{}': must start with http:// or https://",
                    url
                ));
            }
        }
        Ok(())
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // ------------------------------------------------------------------
    // Helper: build a small synthetic ImageRecord
    // ------------------------------------------------------------------

    fn make_record(w: usize, h: usize, label: Option<usize>) -> ImageRecord {
        let pixels: Vec<f32> = (0..h * w * 3).map(|i| (i % 256) as f32 / 255.0).collect();
        ImageRecord {
            pixels,
            width: w,
            height: h,
            channels: 3,
            label,
            key: "test".into(),
        }
    }

    fn make_pipeline(
        n_shards: usize,
        imgs_per_shard: usize,
        batch_size: usize,
    ) -> JepaDataPipeline {
        let shards: Vec<InMemoryShard> = (0..n_shards)
            .map(|id| InMemoryShard::synthetic(imgs_per_shard, 32, id))
            .collect();
        let config = JepaDataConfig {
            image_size: 32,
            patch_size: 8,
            batch_size,
            num_workers: 1,
            shard_paths: Vec::new(),
            scale_range: (0.5, 1.0),
            ratio_range: (0.75, 1.33),
            use_horizontal_flip: true,
            imagenet_normalize: true,
        };
        JepaDataPipeline::new(config, shards)
    }

    // ── MultiScaleRandomCrop ──────────────────────────────────────────────────

    #[test]
    fn test_crop_output_size() {
        let crop = MultiScaleRandomCrop::new(64, (0.2, 1.0), (0.75, 1.33));
        let img = make_record(128, 128, None);
        let out = crop.crop_and_resize(&img, 42);
        assert_eq!(out.len(), 64 * 64 * 3, "output must be [target, target, 3]");
    }

    #[test]
    fn test_crop_output_in_range() {
        // Input pixels are in [0,1], bilinear interp keeps them in [0,1]
        let crop = MultiScaleRandomCrop::new(32, (0.2, 1.0), (0.75, 1.33));
        let img = make_record(64, 64, None);
        let out = crop.crop_and_resize(&img, 7);
        for &v in &out {
            assert!(v >= -1e-6 && v <= 1.0 + 1e-6, "pixel {v} out of [0,1]");
        }
    }

    #[test]
    fn test_crop_different_seeds_differ() {
        let crop = MultiScaleRandomCrop::new(32, (0.2, 0.8), (0.75, 1.33));
        let img = make_record(64, 64, None);
        let a = crop.crop_and_resize(&img, 1);
        let b = crop.crop_and_resize(&img, 999);
        assert_ne!(a, b, "different seeds should produce different crops");
    }

    #[test]
    fn test_crop_same_seed_reproducible() {
        let crop = MultiScaleRandomCrop::new(32, (0.2, 1.0), (0.75, 1.33));
        let img = make_record(64, 64, None);
        let a = crop.crop_and_resize(&img, 42);
        let b = crop.crop_and_resize(&img, 42);
        assert_eq!(a, b, "same seed must be deterministic");
    }

    #[test]
    fn test_crop_small_image() {
        // Should not panic on a 1×1 image
        let crop = MultiScaleRandomCrop::new(8, (0.5, 1.0), (1.0, 1.0));
        let img = make_record(1, 1, None);
        let out = crop.crop_and_resize(&img, 0);
        assert_eq!(out.len(), 8 * 8 * 3);
    }

    // ── RandomHorizontalFlip ──────────────────────────────────────────────────

    #[test]
    fn test_flip_is_own_inverse() {
        let flip = RandomHorizontalFlip::new(1.0); // always flip
        let pixels: Vec<f32> = (0..4 * 4 * 3).map(|i| i as f32 / 48.0).collect();
        let flipped = flip.apply(&pixels, 4, 4, 3, 0);
        let restored = flip.apply(&flipped, 4, 4, 3, 0);
        for (a, b) in pixels.iter().zip(restored.iter()) {
            assert!((a - b).abs() < 1e-6, "double flip should be identity");
        }
    }

    #[test]
    fn test_flip_prob_zero_never_flips() {
        let flip = RandomHorizontalFlip::new(0.0);
        let pixels: Vec<f32> = (0..6 * 6 * 3).map(|i| i as f32 / 100.0).collect();
        for seed in 0..50u64 {
            let out = flip.apply(&pixels, 6, 6, 3, seed);
            assert_eq!(out, pixels, "prob=0 must never flip");
        }
    }

    #[test]
    fn test_flip_prob_one_always_flips() {
        let flip = RandomHorizontalFlip::new(1.0);
        // Use an asymmetric image so flipping is detectable.
        let mut pixels = vec![0.0f32; 4 * 4 * 3];
        pixels[0] = 1.0; // top-left pixel, channel 0
        for seed in 0..10u64 {
            let out = flip.apply(&pixels, 4, 4, 3, seed);
            // After flip, top-left is now what was top-right (which is 0.0)
            assert!((out[0] - 0.0).abs() < 1e-6, "prob=1 must always flip");
            // The original top-left value (1.0) is now at the rightmost position of row 0
            assert!((out[(4 - 1) * 3] - 1.0).abs() < 1e-6);
        }
    }

    #[test]
    fn test_flip_output_same_size() {
        let flip = RandomHorizontalFlip::new(0.5);
        let pixels: Vec<f32> = vec![0.5; 8 * 8 * 3];
        let out = flip.apply(&pixels, 8, 8, 3, 12345);
        assert_eq!(out.len(), pixels.len());
    }

    // ── JepaAugmentationPipeline ──────────────────────────────────────────────

    #[test]
    fn test_normalize_changes_values() {
        let mut pixels = vec![0.5f32; 6];
        let mean = [0.485, 0.456, 0.406];
        let std = [0.229, 0.224, 0.225];
        JepaAugmentationPipeline::normalize(&mut pixels, mean, std);
        // 0.5 normalized should differ from 0.5
        assert!(
            (pixels[0] - 0.5).abs() > 1e-4,
            "normalization must change values"
        );
    }

    #[test]
    fn test_augmentation_pipeline_output_shape() {
        let config = JepaDataConfig {
            image_size: 32,
            patch_size: 8,
            imagenet_normalize: true,
            use_horizontal_flip: true,
            scale_range: (0.5, 1.0),
            ratio_range: (0.75, 1.33),
            ..Default::default()
        };
        let pipeline = JepaAugmentationPipeline::from_config(&config);
        let img = make_record(64, 64, None);
        let out = pipeline.process(&img, 0);
        assert_eq!(out.len(), 32 * 32 * 3);
    }

    #[test]
    fn test_augmentation_consistent_shape_across_samples() {
        let config = JepaDataConfig {
            image_size: 32,
            patch_size: 8,
            imagenet_normalize: true,
            use_horizontal_flip: true,
            scale_range: (0.5, 1.0),
            ratio_range: (0.75, 1.33),
            ..Default::default()
        };
        let pipeline = JepaAugmentationPipeline::from_config(&config);
        for seed in 0..20u64 {
            let img = make_record(64, 64, None);
            let out = pipeline.process(&img, seed);
            assert_eq!(out.len(), 32 * 32 * 3, "seed {seed}: wrong output size");
        }
    }

    // ── InMemoryShard ─────────────────────────────────────────────────────────

    #[test]
    fn test_in_memory_shard_record_count() {
        let shard = InMemoryShard::synthetic(50, 32, 0);
        assert_eq!(shard.len(), 50);
        assert!(!shard.is_empty());
    }

    #[test]
    fn test_in_memory_shard_zero() {
        let shard = InMemoryShard::synthetic(0, 32, 0);
        assert!(shard.is_empty());
    }

    #[test]
    fn test_in_memory_shard_pixel_range() {
        let shard = InMemoryShard::synthetic(5, 16, 1);
        for rec in &shard.records {
            for &px in &rec.pixels {
                assert!(px >= 0.0 && px <= 1.0, "pixel {px} out of [0,1]");
            }
        }
    }

    #[test]
    fn test_in_memory_shard_sequential_labels() {
        let shard = InMemoryShard::synthetic(10, 16, 0);
        for (i, rec) in shard.records.iter().enumerate() {
            assert_eq!(rec.label, Some(i % 1000));
        }
    }

    // ── JepaDataPipeline ─────────────────────────────────────────────────────

    #[test]
    fn test_pipeline_num_batches_per_epoch() {
        // 2 shards × 10 images = 20 total; batch_size=8 → ceil(20/8)=3
        let pipeline = make_pipeline(2, 10, 8);
        assert_eq!(pipeline.num_batches_per_epoch(), 3);
    }

    #[test]
    fn test_pipeline_next_batch_exhausts_epoch() {
        let mut pipeline = make_pipeline(1, 5, 2);
        // 5 images, batch_size=2 → 3 batches (2+2+1)
        let mut count = 0;
        while pipeline.next_batch().is_some() {
            count += 1;
        }
        assert_eq!(count, 3);
        // After exhaustion, next_batch returns None again.
        assert!(pipeline.next_batch().is_none());
    }

    #[test]
    fn test_pipeline_batch_sizes_correct() {
        let mut pipeline = make_pipeline(1, 7, 3);
        // batch 0: 3, batch 1: 3, batch 2: 1
        let b0 = pipeline.next_batch().expect("batch 0");
        assert_eq!(b0.batch_size(), 3);
        let b1 = pipeline.next_batch().expect("batch 1");
        assert_eq!(b1.batch_size(), 3);
        let b2 = pipeline.next_batch().expect("batch 2 (partial)");
        assert_eq!(b2.batch_size(), 1);
        assert!(pipeline.next_batch().is_none());
    }

    #[test]
    fn test_pipeline_reset_epoch_allows_reiteration() {
        let mut pipeline = make_pipeline(1, 4, 4);
        assert!(pipeline.next_batch().is_some());
        assert!(pipeline.next_batch().is_none());
        pipeline.reset_epoch();
        assert_eq!(pipeline.epoch(), 1);
        assert!(pipeline.next_batch().is_some());
    }

    #[test]
    fn test_pipeline_epoch_counter() {
        let mut pipeline = make_pipeline(1, 2, 2);
        assert_eq!(pipeline.epoch(), 0);
        pipeline.reset_epoch();
        assert_eq!(pipeline.epoch(), 1);
        pipeline.reset_epoch();
        assert_eq!(pipeline.epoch(), 2);
    }

    #[test]
    fn test_pipeline_multi_shard() {
        // 3 shards × 4 images = 12 total; batch_size=5 → 3 batches
        let mut pipeline = make_pipeline(3, 4, 5);
        let mut total = 0;
        while let Some(b) = pipeline.next_batch() {
            total += b.batch_size();
        }
        assert_eq!(total, 12);
    }

    #[test]
    fn test_pipeline_batch_size_gt_shard_size() {
        // batch_size larger than any single shard — should still work
        let mut pipeline = make_pipeline(2, 3, 100);
        // 6 total images → 1 batch of 6
        let batch = pipeline.next_batch().expect("should yield one batch");
        assert_eq!(batch.batch_size(), 6);
        assert!(pipeline.next_batch().is_none());
    }

    #[test]
    fn test_pipeline_single_shard() {
        let mut pipeline = make_pipeline(1, 8, 4);
        let b0 = pipeline.next_batch().unwrap();
        let b1 = pipeline.next_batch().unwrap();
        assert_eq!(b0.batch_size() + b1.batch_size(), 8);
        assert!(pipeline.next_batch().is_none());
    }

    // ── JepaBatch mask indices ────────────────────────────────────────────────

    #[test]
    fn test_batch_context_target_non_overlapping() {
        let mut pipeline = make_pipeline(1, 8, 4);
        let batch = pipeline.next_batch().unwrap();
        for (ctx, tgt) in batch
            .context_patch_indices
            .iter()
            .zip(batch.all_target_indices.iter())
        {
            for &ci in ctx {
                assert!(!tgt.contains(&ci), "context index {ci} in target set");
            }
        }
    }

    #[test]
    fn test_batch_target_indices_in_range() {
        let config = JepaDataConfig {
            image_size: 32,
            patch_size: 8,
            batch_size: 4,
            ..JepaDataConfig::default()
        };
        let total_patches = config.num_patches(); // (32/8)^2 = 16
        let shards = vec![InMemoryShard::synthetic(8, 32, 0)];
        let mut pipeline = JepaDataPipeline::new(config, shards);
        let batch = pipeline.next_batch().unwrap();
        for tgt in &batch.all_target_indices {
            for &idx in tgt {
                assert!(idx < total_patches, "target index {idx} >= {total_patches}");
            }
        }
    }

    #[test]
    fn test_batch_context_sorted_deduped() {
        let mut pipeline = make_pipeline(1, 8, 4);
        let batch = pipeline.next_batch().unwrap();
        for ctx in &batch.context_patch_indices {
            for w in ctx.windows(2) {
                assert!(
                    w[0] < w[1],
                    "context indices must be sorted and deduplicated"
                );
            }
        }
    }

    #[test]
    fn test_batch_labels_match_shard() {
        let mut pipeline = make_pipeline(1, 6, 6);
        let batch = pipeline.next_batch().unwrap();
        // All labels must be Some (synthetic shards always set labels)
        for label in &batch.labels {
            assert!(label.is_some());
        }
    }

    #[test]
    fn test_batch_id_increments() {
        let mut pipeline = make_pipeline(1, 9, 3);
        let b0 = pipeline.next_batch().unwrap();
        let b1 = pipeline.next_batch().unwrap();
        let b2 = pipeline.next_batch().unwrap();
        assert_eq!(b0.batch_id, 0);
        assert_eq!(b1.batch_id, 1);
        assert_eq!(b2.batch_id, 2);
    }

    // ── DatasetStats ──────────────────────────────────────────────────────────

    #[test]
    fn test_dataset_stats_totals() {
        let pipeline = make_pipeline(3, 10, 8);
        let stats = DatasetStats::compute(&pipeline);
        assert_eq!(stats.total_images, 30);
        assert_eq!(stats.num_shards, 3);
        assert_eq!(stats.images_per_shard, vec![10, 10, 10]);
    }

    #[test]
    fn test_dataset_stats_mask_efficiency_positive() {
        let pipeline = make_pipeline(1, 8, 4);
        let stats = DatasetStats::compute(&pipeline);
        assert!(stats.mask_efficiency > 0.0, "mask_efficiency should be > 0");
        assert!(
            stats.mask_efficiency <= 1.0,
            "mask_efficiency should be <= 1"
        );
    }

    #[test]
    fn test_dataset_stats_avg_patches_positive() {
        let pipeline = make_pipeline(1, 8, 4);
        let stats = DatasetStats::compute(&pipeline);
        assert!(stats.avg_context_patches > 0.0);
        assert!(stats.avg_target_patches > 0.0);
    }

    // ── WebDatasetShard ───────────────────────────────────────────────────────

    #[test]
    fn test_web_dataset_shard_to_in_memory() {
        let shard = WebDatasetShard::new("/data/imagenet/shard-0000.tar", 100, 0);
        let mem = shard.to_in_memory(32);
        assert_eq!(mem.len(), 100);
        assert_eq!(mem.shard_id, 0);
    }

    #[test]
    fn test_web_dataset_shard_compressed_flag() {
        let gz = WebDatasetShard::new("data/shard.tar.gz", 10, 0);
        assert!(gz.compressed);
        let plain = WebDatasetShard::new("data/shard.tar", 10, 0);
        assert!(!plain.compressed);
    }

    #[test]
    fn test_web_dataset_shard_pixel_count() {
        let shard = WebDatasetShard::new("test.tar", 5, 1);
        let mem = shard.to_in_memory(16);
        for rec in &mem.records {
            assert_eq!(rec.pixels.len(), 16 * 16 * 3);
        }
    }

    // ── Tar parsing helpers ───────────────────────────────────────────────────

    /// Build a tar header block for a regular file with the given filename and size.
    fn make_tar_header(filename: &str, size: u64) -> [u8; 512] {
        let mut block = [0u8; 512];
        // filename: bytes 0..100
        let fname_bytes = filename.as_bytes();
        block[..fname_bytes.len().min(99)]
            .copy_from_slice(&fname_bytes[..fname_bytes.len().min(99)]);
        // size: bytes 124..136, octal ASCII with trailing null
        let octal = format!("{:011o}\0", size);
        block[124..136].copy_from_slice(octal.as_bytes());
        // typeflag: byte 156 = '0' for regular file
        block[156] = b'0';
        // checksum: bytes 148..156, sum of all header bytes with checksum field as spaces
        for i in 148..156 {
            block[i] = b' ';
        }
        let checksum: u32 = block.iter().map(|&b| b as u32).sum();
        let chk_str = format!("{:06o}\0 ", checksum);
        block[148..156].copy_from_slice(chk_str.as_bytes());
        block
    }

    /// Build a tar header block for a directory entry.
    fn make_tar_dir_header(dirname: &str) -> [u8; 512] {
        let mut block = [0u8; 512];
        let fname_bytes = dirname.as_bytes();
        block[..fname_bytes.len().min(99)]
            .copy_from_slice(&fname_bytes[..fname_bytes.len().min(99)]);
        // size = 0
        let octal = format!("{:011o}\0", 0u64);
        block[124..136].copy_from_slice(octal.as_bytes());
        // typeflag '5' = directory
        block[156] = b'5';
        for i in 148..156 {
            block[i] = b' ';
        }
        let checksum: u32 = block.iter().map(|&b| b as u32).sum();
        let chk_str = format!("{:06o}\0 ", checksum);
        block[148..156].copy_from_slice(chk_str.as_bytes());
        block
    }

    /// Build a minimal valid tar archive from (key, image_bytes, label) tuples.
    ///
    /// Each entry emits a `{key}.jpg` file and, if label is Some, a `{key}.cls` file.
    fn make_test_tar(records: &[(&str, &[u8], Option<usize>)]) -> Vec<u8> {
        let mut out: Vec<u8> = Vec::new();

        for (key, image_bytes, label) in records {
            // Image file
            let img_name = format!("{}.jpg", key);
            let img_size = image_bytes.len() as u64;
            out.extend_from_slice(&make_tar_header(&img_name, img_size));
            out.extend_from_slice(image_bytes);
            // Pad to next 512-byte boundary
            let pad = (512 - (image_bytes.len() % 512)) % 512;
            out.extend(std::iter::repeat(0u8).take(pad));

            // Class label file (optional)
            if let Some(lbl) = label {
                let cls_content = format!("{}\n", lbl);
                let cls_bytes = cls_content.as_bytes();
                let cls_name = format!("{}.cls", key);
                out.extend_from_slice(&make_tar_header(&cls_name, cls_bytes.len() as u64));
                out.extend_from_slice(cls_bytes);
                let cls_pad = (512 - (cls_bytes.len() % 512)) % 512;
                out.extend(std::iter::repeat(0u8).take(cls_pad));
            }
        }

        // End-of-archive: two zero blocks
        out.extend(std::iter::repeat(0u8).take(1024));
        out
    }

    // ── Tar parsing tests ─────────────────────────────────────────────────────

    #[test]
    fn test_webdataset_shard_load_reads_real_tar() {
        // Write a real tar to a temp file and load it through WebDatasetShard.
        let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
        let tar = make_test_tar(&[("000000", fake_jpg, Some(7)), ("000001", fake_jpg, None)]);
        let dir = std::env::temp_dir();
        let path = dir.join(format!("jepa_shard_load_test_{}.tar", std::process::id()));
        std::fs::write(&path, &tar).expect("write temp tar");

        let shard = WebDatasetShard::new(path.to_str().unwrap(), 2, 3);
        let mem = shard.load().expect("load real tar");
        std::fs::remove_file(&path).ok();

        assert_eq!(mem.shard_id, 3);
        assert_eq!(mem.records.len(), 2);
        assert_eq!(mem.records[0].label, Some(7));
        assert_eq!(mem.records[1].label, None);
    }

    #[test]
    fn test_webdataset_shard_load_reads_gzip_tar() {
        use std::io::Write;

        let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
        let tar = make_test_tar(&[("000000", fake_jpg, Some(11))]);

        let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
        encoder.write_all(&tar).expect("gzip write");
        let gz = encoder.finish().expect("gzip finish");

        let dir = std::env::temp_dir();
        let path = dir.join(format!("jepa_shard_gz_test_{}.tar.gz", std::process::id()));
        std::fs::write(&path, &gz).expect("write temp tar.gz");

        let shard = WebDatasetShard::new(path.to_str().unwrap(), 1, 5);
        let mem = shard.load().expect("load gzip tar");
        std::fs::remove_file(&path).ok();

        assert!(shard.compressed);
        assert_eq!(mem.shard_id, 5);
        assert_eq!(mem.records.len(), 1);
        assert_eq!(mem.records[0].label, Some(11));
    }

    #[test]
    fn test_parse_empty_tar() {
        // Two zero blocks = empty archive
        let data = vec![0u8; 1024];
        let records = parse_tar_bytes(&data);
        assert_eq!(records.len(), 0, "empty tar must yield 0 records");
    }

    #[test]
    fn test_parse_single_record() {
        let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
        let tar = make_test_tar(&[("000000", fake_jpg, Some(42))]);
        let records = parse_tar_bytes(&tar);
        assert_eq!(records.len(), 1);
        assert_eq!(records[0].key, "000000");
        assert_eq!(records[0].label, Some(42));
        assert_eq!(records[0].extension, "jpg");
        assert_eq!(records[0].image_bytes, fake_jpg);
    }

    #[test]
    fn test_parse_multiple_records() {
        let fake_jpg = b"\xFF\xD8content";
        let entries: Vec<(&str, &[u8], Option<usize>)> = vec![
            ("000000", fake_jpg, Some(0)),
            ("000001", fake_jpg, Some(1)),
            ("000002", fake_jpg, Some(2)),
            ("000003", fake_jpg, Some(3)),
            ("000004", fake_jpg, Some(4)),
        ];
        let tar = make_test_tar(&entries);
        let records = parse_tar_bytes(&tar);
        assert_eq!(records.len(), 5, "expected 5 records");
        for (i, rec) in records.iter().enumerate() {
            assert_eq!(rec.label, Some(i));
        }
    }

    #[test]
    fn test_parse_record_without_cls() {
        let fake_jpg = b"\xFF\xD8no label";
        let tar = make_test_tar(&[("000000", fake_jpg, None)]);
        let records = parse_tar_bytes(&tar);
        assert_eq!(records.len(), 1);
        assert_eq!(
            records[0].label, None,
            "image with no .cls should have label=None"
        );
    }

    #[test]
    fn test_parse_directory_entries_skipped() {
        let mut out: Vec<u8> = Vec::new();
        // Directory entry
        out.extend_from_slice(&make_tar_dir_header("somedir/"));
        // A regular image record after the directory
        let fake_jpg = b"\xFF\xD8data";
        let img_name = "000000.jpg";
        out.extend_from_slice(&make_tar_header(img_name, fake_jpg.len() as u64));
        out.extend_from_slice(fake_jpg);
        let pad = (512 - (fake_jpg.len() % 512)) % 512;
        out.extend(std::iter::repeat(0u8).take(pad));
        // End-of-archive
        out.extend(std::iter::repeat(0u8).take(1024));

        let records = parse_tar_bytes(&out);
        assert_eq!(records.len(), 1, "directory entries must be skipped");
        assert_eq!(records[0].key, "000000");
    }

    #[test]
    fn test_parse_large_file() {
        // 50 000 bytes > 512*97 = 49664, spans many blocks
        let big_image: Vec<u8> = (0..50_000u32).map(|i| (i % 251) as u8).collect();
        let tar = make_test_tar(&[("000000", &big_image, Some(7))]);
        let records = parse_tar_bytes(&tar);
        assert_eq!(records.len(), 1);
        assert_eq!(records[0].image_bytes.len(), 50_000);
        assert_eq!(records[0].image_bytes, big_image);
    }

    // ── ShuffleBuffer tests ───────────────────────────────────────────────────

    fn make_image_record(key: &str) -> ImageRecord {
        ImageRecord {
            pixels: vec![0.5f32; 8 * 8 * 3],
            width: 8,
            height: 8,
            channels: 3,
            label: None,
            key: key.to_owned(),
        }
    }

    #[test]
    fn test_shuffle_buffer_fill() {
        let mut buf = ShuffleBuffer::new(10);
        for i in 0..10 {
            buf.push(make_image_record(&format!("img{i}")));
        }
        assert_eq!(buf.len(), 10);
        assert!(buf.is_full());
    }

    #[test]
    fn test_shuffle_buffer_push_overflow() {
        let mut buf = ShuffleBuffer::new(10);
        for i in 0..15 {
            buf.push(make_image_record(&format!("img{i}")));
        }
        assert_eq!(buf.len(), 10, "buffer must not grow beyond capacity");
    }

    #[test]
    fn test_shuffle_buffer_drain() {
        let mut buf = ShuffleBuffer::new(10);
        for i in 0..10 {
            buf.push(make_image_record(&format!("img{i}")));
        }
        let batch = buf.drain_batch(5);
        assert_eq!(batch.len(), 5);
        assert_eq!(buf.len(), 5);
    }

    #[test]
    fn test_shuffle_buffer_empty_drain() {
        let mut buf = ShuffleBuffer::new(10);
        for i in 0..3 {
            buf.push(make_image_record(&format!("img{i}")));
        }
        // Drain more than available
        let batch = buf.drain_batch(100);
        assert_eq!(
            batch.len(),
            3,
            "drain should return at most what's in the buffer"
        );
        assert_eq!(buf.len(), 0);

        // Drain zero
        let empty = buf.drain_batch(0);
        assert_eq!(empty.len(), 0);
    }

    // ── ShardLoadStats tests ──────────────────────────────────────────────────

    #[test]
    fn test_shard_load_stats_bytes() {
        use std::io::Write;

        // Write a real tar to a temp file and read it back.
        let fake_jpg = b"\xFF\xD8fake";
        let entries: Vec<(&str, &[u8], Option<usize>)> = vec![
            ("000000", fake_jpg, Some(0)),
            ("000001", fake_jpg, Some(1)),
            ("000002", fake_jpg, Some(2)),
        ];
        let tar_bytes = make_test_tar(&entries);

        let mut tmp = std::env::temp_dir();
        tmp.push("jepa_test_shard_stats.tar");
        {
            let mut f = std::fs::File::create(&tmp).expect("create temp file");
            f.write_all(&tar_bytes).expect("write tar");
        }

        let (records, stats) = read_webdataset_shard(&tmp).expect("read shard");
        std::fs::remove_file(&tmp).ok();

        assert_eq!(records.len(), 3);
        assert!(stats.bytes_read > 0, "bytes_read must be > 0");
        assert_eq!(stats.records_loaded, 3);
    }

    // ── from_filesystem tests ─────────────────────────────────────────────────

    #[test]
    fn test_from_filesystem_missing_path() {
        let config = JepaDataConfig::default();
        let missing = std::path::PathBuf::from("/nonexistent/path/shard-0000.tar");
        let result = JepaDataPipeline::from_filesystem(config, vec![missing.clone()]);
        assert!(result.is_err(), "missing path must return Err");
        // Extract the error message without calling unwrap_err (JepaDataPipeline has no Debug).
        let msg = match result {
            Err(e) => e,
            Ok(_) => panic!("expected Err"),
        };
        assert!(
            msg.contains("nonexistent") || msg.contains("shard-0000.tar"),
            "error message should mention the path: {msg}"
        );
    }

    // ── webdataset_record_to_image ────────────────────────────────────────────

    #[test]
    fn test_webdataset_decode_placeholder_shape() {
        let rec = WebDatasetRecord {
            key: "test".to_string(),
            image_bytes: vec![0u8; 100], // not a valid image
            label: Some(1),
            extension: "jpg".to_string(),
        };
        let img = webdataset_record_to_image(rec);
        assert_eq!(img.pixels.len(), 224 * 224 * 3);
        assert_eq!(img.width, 224);
        assert_eq!(img.height, 224);
        assert_eq!(img.channels, 3);
    }

    #[test]
    fn test_webdataset_decode_placeholder_value() {
        let rec = WebDatasetRecord {
            key: "k".to_string(),
            image_bytes: vec![],
            label: None,
            extension: "jpg".to_string(),
        };
        let img = webdataset_record_to_image(rec);
        for &p in &img.pixels {
            assert!(
                (p - 0.5f32).abs() < 1e-6 || (0.0..=1.0).contains(&p),
                "pixel {} must be in [0, 1]",
                p
            );
        }
    }

    #[test]
    fn test_webdataset_decode_label_preserved() {
        let rec = WebDatasetRecord {
            key: "sample_0000".to_string(),
            image_bytes: vec![0xFF, 0xD8, 0xFF], // JPEG magic start
            label: Some(42),
            extension: "jpg".to_string(),
        };
        let img = webdataset_record_to_image(rec);
        assert_eq!(img.label, Some(42));
    }

    #[test]
    fn test_webdataset_decode_key_preserved() {
        let rec = WebDatasetRecord {
            key: "my_key_123".to_string(),
            image_bytes: vec![],
            label: None,
            extension: "png".to_string(),
        };
        let img = webdataset_record_to_image(rec);
        assert_eq!(img.key, "my_key_123");
    }

    // ── HTTP WebDataset ──────────────────────────────────────────────────────

    // test: validate_urls accepts valid http URLs
    #[test]
    fn test_validate_urls_http_ok() {
        let urls = vec![
            "http://example.com/shard-0.tar".to_string(),
            "https://storage.example.com/data/shard-1.tar".to_string(),
        ];
        assert!(JepaDataPipeline::validate_urls(&urls).is_ok());
    }

    // test: validate_urls rejects non-http schemes
    #[test]
    fn test_validate_urls_bad_scheme() {
        let urls = vec!["s3://bucket/shard.tar".to_string()];
        assert!(JepaDataPipeline::validate_urls(&urls).is_err());
    }

    // test: validate_urls empty list is Ok
    #[test]
    fn test_validate_urls_empty() {
        assert!(JepaDataPipeline::validate_urls(&[]).is_ok());
    }

    // test: from_urls with bad scheme returns Err immediately
    #[test]
    fn test_from_urls_bad_scheme_err() {
        let config = JepaDataConfig::default();
        let result =
            JepaDataPipeline::from_urls(config, vec!["ftp://example.com/shard.tar".to_string()]);
        assert!(result.is_err());
        let msg = match result {
            Err(e) => e,
            Ok(_) => panic!("expected Err"),
        };
        assert!(msg.contains("unsupported URL scheme"), "error was: {msg}");
    }

    // test: download_webdataset_shard with invalid URL scheme returns Err
    #[test]
    fn test_download_bad_scheme() {
        let result = download_webdataset_shard("file:///tmp/test.tar", 0);
        assert!(result.is_err());
    }

    // test: download_webdataset_shard to a non-existent host returns Err
    // 192.0.2.0/24 is reserved "documentation" space — guaranteed unreachable.
    // Uses a 5-second timeout so the test completes quickly.
    #[test]
    fn test_download_unreachable_host() {
        let result =
            download_webdataset_shard_with_timeout("http://192.0.2.1:9999/shard.tar", 0, 5);
        assert!(result.is_err(), "unreachable host should return Err");
    }

    // test: UrlShardDescriptor creation
    #[test]
    fn test_url_shard_descriptor_new() {
        let desc = UrlShardDescriptor::new("https://example.com/shard-0.tar", 1000, 0);
        assert_eq!(desc.expected_records, 1000);
        assert_eq!(desc.shard_id, 0);
        assert!(desc.is_https());
    }

    // test: UrlShardDescriptor::is_https with http URL
    #[test]
    fn test_url_shard_descriptor_is_http() {
        let desc = UrlShardDescriptor::new("http://example.com/shard.tar", 0, 1);
        assert!(!desc.is_https());
    }

    // test: UrlShardDescriptor clone
    #[test]
    fn test_url_shard_descriptor_clone() {
        let desc = UrlShardDescriptor::new("https://a.b/c.tar", 42, 3);
        let desc2 = desc.clone();
        assert_eq!(desc2.url, desc.url);
        assert_eq!(desc2.shard_id, 3);
    }

    // test: from_urls with empty list succeeds with empty pipeline
    #[test]
    fn test_from_urls_empty_list() {
        let config = JepaDataConfig::default();
        // Empty URL list → empty pipeline (no shards loaded)
        let result = JepaDataPipeline::from_urls(config, vec![]);
        assert!(result.is_ok());
    }
}

#[cfg(all(test, feature = "image-decode"))]
mod tests_image_decode {
    use super::*;

    #[test]
    fn test_webdataset_decode_real_image() {
        // Create a minimal 1×1 PNG using the image crate itself
        use image::{ImageBuffer, Rgb};
        let img_buf: ImageBuffer<Rgb<u8>, Vec<u8>> =
            ImageBuffer::from_pixel(1, 1, Rgb([128u8, 64u8, 32u8]));
        let mut bytes = Vec::new();
        img_buf
            .write_to(
                &mut std::io::Cursor::new(&mut bytes),
                image::ImageOutputFormat::Png,
            )
            .unwrap();

        let rec = WebDatasetRecord {
            key: "synthetic".to_string(),
            image_bytes: bytes,
            label: Some(0),
            extension: "png".to_string(),
        };
        let decoded = webdataset_record_to_image(rec);
        assert_eq!(decoded.pixels.len(), 224 * 224 * 3);
        assert_eq!(decoded.width, 224);
        assert_eq!(decoded.height, 224);
        for &p in &decoded.pixels {
            assert!((0.0..=1.0).contains(&p), "pixel {} out of range", p);
        }
    }
}