//! High-level access to an RSSA file: header parsing plus eager and streaming access
//! to the track data, backed by a memory-mapped view of the file.

use std::fs::File;
use std::path::Path;
use std::sync::Arc;

use memmap2::Mmap;
use polars::prelude::DataFrame;

use crate::error::{Result, RssaError};
use crate::header::{self, FileParameters};
use crate::tracks::{self, Column, RECORD_BYTES, TrackLayout};

/// A memory-mapped RSSA file, with its header already parsed.
#[derive(Debug)]
pub struct RssaFile {
    mmap: Arc<Mmap>,
    parameters: FileParameters,
    layout: TrackLayout,
    data_offset: usize,
}

impl RssaFile {
    /// Opens and memory-maps `path`, parses its header, and validates that the
    /// amount of track data present matches what the header declares.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let file = File::open(path)?;
        // SAFETY: memory-mapping a file is unsafe because the mapping becomes
        // invalid (UB) if the file is truncated or its contents change while
        // mapped. RSSA files are treated as immutable simulation outputs that are
        // fully written before being read; we don't attempt to guard against
        // concurrent external modification, which mirrors the standard use of
        // mmap for read-only access to static files.
        let mmap = unsafe { Mmap::map(&file)? };
        let (parameters, data_offset) = header::parse_header(&mmap)?;

        let track_bytes = mmap.len().saturating_sub(data_offset);
        let expected_bytes = parameters.nrss.max(0) as u64 * RECORD_BYTES as u64;
        if track_bytes as u64 != expected_bytes {
            return Err(RssaError::TrackDataSizeMismatch {
                nrss: parameters.nrss,
                expected_bytes,
                actual_bytes: track_bytes as u64,
            });
        }

        Ok(Self {
            mmap: Arc::new(mmap),
            layout: TrackLayout::from_parameters(&parameters),
            parameters,
            data_offset,
        })
    }

    pub fn parameters(&self) -> &FileParameters {
        &self.parameters
    }

    fn track_bytes(&self) -> &[u8] {
        &self.mmap[self.data_offset..]
    }

    /// Decodes track data into a [`DataFrame`] containing `columns`.
    ///
    /// If `max_tracks` is given, only the first `max_tracks` records are decoded
    /// (any records past that point in the file are never read), which keeps a
    /// `.head(n)`-style query on a very large file cheap.
    pub fn read_tracks(&self, columns: &[Column], max_tracks: Option<usize>) -> Result<DataFrame> {
        let bytes = self.track_bytes();
        let bytes = match max_tracks {
            Some(n) => &bytes[..(n * RECORD_BYTES).min(bytes.len())],
            None => bytes,
        };
        tracks::decode_tracks(bytes, columns, &self.layout)
    }

    /// Returns a streaming batch reader over the track data, decoding `batch_size`
    /// tracks at a time.
    pub fn batches(&self, batch_size: usize, columns: Vec<Column>) -> BatchReader {
        BatchReader {
            mmap: Arc::clone(&self.mmap),
            layout: self.layout.clone(),
            data_offset: self.data_offset,
            total_tracks: self.parameters.nrss.max(0) as usize,
            cursor_track: 0,
            batch_size: batch_size.max(1),
            columns,
        }
    }
}

/// Sequential, cursor-based reader that decodes the track data of an [`RssaFile`] in
/// batches, without needing to keep the whole (potentially 100+ GB) decoded result in
/// memory at once.
pub struct BatchReader {
    mmap: Arc<Mmap>,
    layout: TrackLayout,
    data_offset: usize,
    total_tracks: usize,
    cursor_track: usize,
    batch_size: usize,
    columns: Vec<Column>,
}

impl BatchReader {
    /// Decodes and returns the next batch, or `None` once the end of the file has
    /// been reached.
    pub fn next_batch(&mut self) -> Result<Option<DataFrame>> {
        if self.cursor_track >= self.total_tracks {
            return Ok(None);
        }
        let n = self.remaining_tracks().min(self.batch_size);
        let start = self.data_offset + self.cursor_track * RECORD_BYTES;
        let end = start + n * RECORD_BYTES;
        let df = tracks::decode_tracks(&self.mmap[start..end], &self.columns, &self.layout)?;
        self.cursor_track += n;
        Ok(Some(df))
    }

    pub fn remaining_tracks(&self) -> usize {
        self.total_tracks - self.cursor_track
    }
}

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

    fn small_cyl_path() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../resources/small_cyl.w")
    }

    fn multiple_surfs_path() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../resources/multiple_surfs.w")
    }

    #[test]
    fn decodes_the_surface_id_of_each_track_when_the_file_records_one() {
        let file = RssaFile::open(multiple_surfs_path()).unwrap();
        let df = file.read_tracks(&[Column::SurfaceId], None).unwrap();
        let ids = df.column("surface_id").unwrap().i64().unwrap();
        assert_eq!(ids.null_count(), 0);

        let mut counts = std::collections::HashMap::new();
        for id in ids.iter().flatten() {
            *counts.entry(id).or_insert(0usize) += 1;
        }
        // The header writes this plane's id as -610041; the tracks carry it unsigned.
        assert_eq!(counts[&610_041], 32_832);
        assert_eq!(counts[&63], 9_936);
        assert_eq!(counts[&252], 582);
    }

    #[test]
    fn surface_id_is_null_when_the_file_does_not_record_one() {
        // small_cyl.w has nrcd > 0: the last value of each track record is a direction
        // cosine, not a surface id. Truncating it would label every track "surface 0",
        // so the column has to come out empty instead.
        let file = RssaFile::open(small_cyl_path()).unwrap();
        let df = file.read_tracks(&[Column::SurfaceId], None).unwrap();
        let ids = df.column("surface_id").unwrap().i64().unwrap();
        assert_eq!(ids.len(), 72_083);
        assert_eq!(ids.null_count(), 72_083);
    }

    #[test]
    fn opens_and_validates_the_small_cyl_fixture() {
        let file = RssaFile::open(small_cyl_path()).unwrap();
        assert_eq!(file.parameters().nrss, 72_083);
    }

    #[test]
    fn read_tracks_decodes_the_full_file() {
        let file = RssaFile::open(small_cyl_path()).unwrap();
        let df = file.read_tracks(&Column::ALL, None).unwrap();
        assert_eq!(df.shape(), (72_083, 12));

        let neutrons = df
            .column("particle_type")
            .unwrap()
            .i64()
            .unwrap()
            .iter()
            .filter(|v| *v == Some(tracks::NEUTRON_PARTICLE_TYPE))
            .count();
        assert!(neutrons > 0);
        let photons = 72_083 - neutrons;
        assert_eq!(photons, 0);
    }

    #[test]
    fn read_tracks_respects_max_tracks() {
        let file = RssaFile::open(small_cyl_path()).unwrap();
        let full = file.read_tracks(&Column::ALL, None).unwrap();
        let limited = file.read_tracks(&Column::ALL, Some(20)).unwrap();
        assert_eq!(limited.shape(), (20, 12));
        // equals_missing, not equals: small_cyl.w records no per-track surface id, so
        // that column is all-null on both sides and null != null under equals().
        assert!(limited.equals_missing(&full.head(Some(20))));
    }

    #[test]
    fn batch_reader_matches_eager_read() {
        let file = RssaFile::open(small_cyl_path()).unwrap();
        let eager = file.read_tracks(&Column::ALL, None).unwrap();

        let mut reader = file.batches(1_000, Column::ALL.to_vec());
        let mut batches = Vec::new();
        while let Some(batch) = reader.next_batch().unwrap() {
            batches.push(batch);
        }
        assert_eq!(batches.len(), 73); // 72 full batches of 1000 + 1 partial batch

        let mut concatenated = batches.remove(0);
        for batch in batches {
            concatenated.vstack_mut(&batch).unwrap();
        }
        assert!(eager.equals_missing(&concatenated));
    }

    #[test]
    fn rejects_a_truncated_file() {
        let bytes = std::fs::read(small_cyl_path()).unwrap();
        let truncated = &bytes[..bytes.len() - 10];
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("truncated.w");
        std::fs::write(&path, truncated).unwrap();

        let err = RssaFile::open(&path).unwrap_err();
        assert!(matches!(err, RssaError::TrackDataSizeMismatch { .. }));
    }
}
