//! Parsing of the RSSA file header: the file-level parameters and the list of
//! surfaces the tracks were recorded on.

use serde::{Deserialize, Serialize};

use crate::error::{Result, RssaError};
use crate::fortran::{RecordCursor, read_f64, read_i32, read_i64};

/// Number of values recorded for each particle track (history, packed particle type,
/// weight, energy, time, x, y, z, u, v, surface id).
pub const VALUES_PER_TRACK: i32 = 11;

/// Version of the JSON representation of [`FileParameters`], written into every
/// serialized copy and checked when reading one back.
///
/// Version 1 (up to v0.0.6) parsed the surface records with the wrong layout and typed
/// their parameters as integers, so the values it stored are silently wrong rather
/// than merely differently shaped - which is exactly why they must be rejected instead
/// of loaded.
pub const PARAMETERS_JSON_FORMAT_VERSION: u32 = 2;

/// The code that produced the RSSA file, identified from its first header record.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SourceCode {
    /// Vanilla MCNP, identified by a `"SF_00001"` marker.
    Mcnp,
    /// D1S-UNED (a modified MCNP), identified by a `"d1suned"` marker.
    D1sUned,
}

/// Parameters of a single surface tracks were recorded on.
///
/// `surface_type` is the index of the surface's mnemonic in MCNP's standard
/// surface-card order (`p` = 1, `px` = 2, ..., `pz` = 4, ..., `cz` = 15, ...), and
/// `parameters` holds its coefficients as MCNP stores them internally - notably as
/// *squared* radii for cylinders and spheres. Both are interpreted on the Python side
/// (see `aci_rssa.header.SurfaceType`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SurfaceParameters {
    pub id: i32,
    /// Extra macrobody-facet info; only meaningful when `kjaq == 1`, `-1` otherwise.
    pub info: i32,
    pub surface_type: i32,
    pub num_parameters: i32,
    pub parameters: Vec<f64>,
}

/// Parameters extracted from an RSSA file header.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FileParameters {
    /// Version of this struct's JSON representation; see
    /// [`PARAMETERS_JSON_FORMAT_VERSION`].
    pub format_version: u32,
    pub source_code: SourceCode,
    /// Number of histories of the simulation that produced this file, given as a
    /// negative number (as written by MCNP).
    pub np1: i64,
    /// Number of tracks recorded.
    pub nrss: i64,
    /// Number of values recorded per track (always 11, checked at parse time).
    pub nrcd: i32,
    /// Number of surfaces tracks were recorded on.
    pub njsw: i32,
    /// Number of independent histories that reached the recorded surfaces.
    pub niss: i64,
    /// Number of extra (weight-window) records; only files with `niwr == 0` are
    /// supported.
    pub niwr: i32,
    /// Source particle type.
    pub mipts: i32,
    /// Flag indicating whether surfaces include macrobody facet info.
    pub kjaq: i32,
    /// Present only for D1S-UNED files: the index of the last recorded dump.
    pub last_dump: Option<i32>,
    pub surfaces: Vec<SurfaceParameters>,
}

impl FileParameters {
    /// Whether each track record carries the id of the surface it crossed in its last
    /// value.
    ///
    /// MCNP 6 files (`nrcd < 0`, e.g. `resources/multiple_surfs.w`) record it there;
    /// files written with `nrcd > 0` (e.g. the D1S-UNED `resources/small_cyl.w`) put a
    /// direction cosine in that slot instead and encode the surface elsewhere, in a
    /// packed field this crate does not decode. Reading the cosine as an id yields 0
    /// for every track, so the distinction has to be made here rather than silently
    /// producing a column of zeroes.
    pub fn records_surface_id(&self) -> bool {
        self.nrcd < 0
    }

    /// The ids of the surfaces tracks were recorded on, as written in the header.
    ///
    /// Note these can be negative (MCNP writes the surface of `resources/multiple_surfs.w`
    /// as `-610041` while its tracks carry `610041`), so comparisons against a track's
    /// surface id are made on absolute values.
    pub fn surface_ids(&self) -> Vec<i32> {
        self.surfaces.iter().map(|s| s.id).collect()
    }

    pub fn to_json(&self) -> Result<String> {
        Ok(serde_json::to_string_pretty(self)?)
    }

    pub fn from_json(text: &str) -> Result<Self> {
        Ok(serde_json::from_str(text)?)
    }
}

/// Parses the RSSA header starting at byte 0 of `data`.
///
/// Returns the parsed [`FileParameters`] together with the byte offset at which the
/// track data begins.
pub fn parse_header(data: &[u8]) -> Result<(FileParameters, usize)> {
    let mut cursor = RecordCursor::new(data);

    let id_record = cursor.read_record()?;
    let id_text = String::from_utf8_lossy(id_record);
    let (source_code, last_dump) = if id_text.contains("d1suned") {
        let last_dump = if id_record.len() >= 4 {
            Some(read_i32(id_record, id_record.len() - 4)?)
        } else {
            None
        };
        (SourceCode::D1sUned, last_dump)
    } else if id_text.contains("SF_00001") {
        let _code_version = cursor.read_record()?; // informational only, not needed
        (SourceCode::Mcnp, None)
    } else {
        return Err(RssaError::UnsupportedFormat(id_text.into_owned()));
    };

    let dims_record = cursor.read_record()?;
    let np1 = read_i64(dims_record, 0)?;
    let nrss = read_i64(dims_record, 8)?;
    let nrcd = read_i32(dims_record, 16)?;
    let njsw = read_i32(dims_record, 20)?;
    let niss = read_i64(dims_record, 24)?;

    if nrcd.abs() != VALUES_PER_TRACK {
        return Err(RssaError::UnexpectedRecordLength {
            expected: VALUES_PER_TRACK,
            found: nrcd,
        });
    }
    if np1 >= 0 {
        return Err(RssaError::UnsupportedNp1(np1));
    }

    let extra_record = cursor.read_record()?;
    let niwr = read_i32(extra_record, 0)?;
    let mipts = read_i32(extra_record, 4)?;
    let kjaq = read_i32(extra_record, 8)?;

    let mut surfaces = Vec::with_capacity(njsw.max(0) as usize);
    for index in 0..njsw as usize {
        let record = cursor.read_record()?;
        // Layout: id, [macrobody-facet info, only when kjaq == 1], surface type and
        // parameter count as 4-byte integers, followed by the surface's coefficients
        // as 8-byte floats. The facet field really is absent (not merely unused) when
        // kjaq == 0, so everything after it shifts; and the coefficients really are
        // floats (a cylinder's squared radius, a plane's d, ...), so reading them as
        // integers yields garbage. Both were wrong up to v0.0.6.
        let id = read_i32(record, 0)?;
        let mut offset = 4;
        let info = if kjaq == 1 {
            let info = read_i32(record, offset)?;
            offset += 4;
            info
        } else {
            -1
        };
        let surface_type = read_i32(record, offset)?;
        let num_parameters = read_i32(record, offset + 4)?;
        offset += 8;

        // Check the record's length against the declared parameter count *before*
        // allocating: a misread count (which is what the pre-v0.0.7 layout produced,
        // e.g. 1_546_188_228 from a float's bytes) would otherwise reserve gigabytes
        // before failing.
        let expected = offset + 8 * num_parameters.max(0) as usize;
        if num_parameters < 0 || expected != record.len() {
            return Err(RssaError::UnexpectedSurfaceRecordLength {
                index,
                num_parameters,
                expected,
                found: record.len(),
                info: if kjaq == 1 {
                    " and the macrobody-facet field"
                } else {
                    ""
                },
            });
        }

        let mut parameters = Vec::with_capacity(num_parameters as usize);
        for i in 0..num_parameters as usize {
            parameters.push(read_f64(record, offset + 8 * i)?);
        }
        surfaces.push(SurfaceParameters {
            id,
            info,
            surface_type,
            num_parameters,
            parameters,
        });
    }

    if niwr != 0 {
        return Err(RssaError::UnsupportedNiwr(niwr));
    }

    let _summary_record = cursor.read_record()?; // contents currently unused

    let parameters = FileParameters {
        format_version: PARAMETERS_JSON_FORMAT_VERSION,
        source_code,
        np1,
        nrss,
        nrcd,
        njsw,
        niss,
        niwr,
        mipts,
        kjaq,
        last_dump,
        surfaces,
    };
    Ok((parameters, cursor.position()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fortran::test_support::write_record;
    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 parses_the_small_cyl_fixture() {
        let bytes = std::fs::read(small_cyl_path()).expect("resources/small_cyl.w must exist");
        let (params, _offset) = parse_header(&bytes).unwrap();

        assert_eq!(params.source_code, SourceCode::D1sUned);
        assert_eq!(params.np1, -100_000);
        assert_eq!(params.nrss, 72_083);
        assert_eq!(params.nrcd, 11);
        assert_eq!(params.njsw, 1);
        assert_eq!(params.niss, 70_797);
        assert_eq!(params.niwr, 0);
        assert_eq!(params.mipts, 3);
        assert_eq!(params.kjaq, 0);

        assert_eq!(params.surfaces.len(), 1);
        assert_eq!(params.surfaces[0].id, 1);
        assert_eq!(params.surfaces[0].info, -1);
        // A `cz` cylinder of radius 30 cm: MCNP stores the *squared* radius. Every
        // track in this file sits at hypot(x, y) == 30.
        assert_eq!(params.surfaces[0].surface_type, 15);
        assert_eq!(params.surfaces[0].num_parameters, 1);
        assert_eq!(params.surfaces[0].parameters, vec![900.0]);
        assert!(!params.records_surface_id()); // nrcd > 0: no per-track surface id
    }

    #[test]
    fn parses_the_multiple_surfs_fixture() {
        let bytes = std::fs::read(multiple_surfs_path()).expect("multiple_surfs.w must exist");
        let (params, _offset) = parse_header(&bytes).unwrap();

        assert_eq!(params.source_code, SourceCode::Mcnp);
        assert_eq!(params.np1, -100_000);
        assert_eq!(params.nrss, 43_350);
        assert_eq!(params.nrcd, -11);
        assert_eq!(params.njsw, 3);
        assert_eq!(params.kjaq, 0);
        assert!(params.records_surface_id());

        // The three surfaces, cross-checked against the tracks recorded on each of
        // them: a plane at z = -800, a cylinder of radius sqrt(26994.49) = 164.3 about
        // the z axis, and a tilted general plane.
        assert_eq!(params.surface_ids(), vec![-610_041, 63, 252]);

        assert_eq!(params.surfaces[0].surface_type, 4); // pz
        assert_eq!(params.surfaces[0].parameters, vec![-800.0]);

        assert_eq!(params.surfaces[1].surface_type, 15); // cz
        assert!((params.surfaces[1].parameters[0] - 26_994.49).abs() < 1e-9);

        assert_eq!(params.surfaces[2].surface_type, 1); // p
        assert_eq!(params.surfaces[2].num_parameters, 4);
        assert_eq!(params.surfaces[2].parameters[3], -34.75);
    }

    #[test]
    fn json_round_trip() {
        let bytes = std::fs::read(small_cyl_path()).expect("resources/small_cyl.w must exist");
        let (params, _offset) = parse_header(&bytes).unwrap();
        let json = params.to_json().unwrap();
        let round_tripped = FileParameters::from_json(&json).unwrap();
        assert_eq!(params, round_tripped);
    }

    /// Builds a synthetic MCNP-style (non D1S-UNED) header with the given surfaces, so
    /// branches unreachable via `small_cyl.w` (multiple surfaces, the `SF_00001`
    /// identification variant, `kjaq == 1`) can still be exercised.
    fn synthetic_mcnp_header(
        njsw: i32,
        niwr: i32,
        kjaq: i32,
        surfaces: &[(i32, i32, i32, &[f64])],
    ) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend(write_record(
            b"SF_00001 the rest of the id record is free-form text",
        ));
        out.extend(write_record(b"code version info, unused"));

        let mut dims = Vec::new();
        dims.extend_from_slice(&(-100i64).to_le_bytes()); // np1 < 0
        dims.extend_from_slice(&10i64.to_le_bytes()); // nrss
        dims.extend_from_slice(&11i32.to_le_bytes()); // nrcd
        dims.extend_from_slice(&njsw.to_le_bytes());
        dims.extend_from_slice(&5i64.to_le_bytes()); // niss
        out.extend(write_record(&dims));

        let mut extra = Vec::new();
        extra.extend_from_slice(&niwr.to_le_bytes());
        extra.extend_from_slice(&3i32.to_le_bytes()); // mipts
        extra.extend_from_slice(&kjaq.to_le_bytes());
        out.extend(write_record(&extra));

        for &(id, info, surface_type, params) in surfaces {
            let mut record = Vec::new();
            record.extend_from_slice(&id.to_le_bytes());
            if kjaq == 1 {
                record.extend_from_slice(&info.to_le_bytes());
            }
            record.extend_from_slice(&surface_type.to_le_bytes());
            record.extend_from_slice(&(params.len() as i32).to_le_bytes());
            for &p in params {
                record.extend_from_slice(&p.to_le_bytes());
            }
            out.extend(write_record(&record));
        }

        out.extend(write_record(b"summary record, unused"));
        out
    }

    #[test]
    fn parses_mcnp_variant_with_multiple_surfaces_and_macrobody_info() {
        let surfaces = [
            (1, 42, 1, [].as_slice()),
            (2, -7, 3, [10.5, -20.25, 30.0].as_slice()),
        ];
        let bytes = synthetic_mcnp_header(2, 0, 1, &surfaces);
        let (params, offset) = parse_header(&bytes).unwrap();

        assert_eq!(params.source_code, SourceCode::Mcnp);
        assert_eq!(params.last_dump, None);
        assert_eq!(params.njsw, 2);
        assert_eq!(params.surfaces.len(), 2);
        assert_eq!(params.surfaces[0].info, 42);
        assert_eq!(params.surfaces[1].id, 2);
        assert_eq!(params.surfaces[1].parameters, vec![10.5, -20.25, 30.0]);
        assert_eq!(offset, bytes.len());
    }

    #[test]
    fn kjaq_zero_omits_the_info_field_entirely() {
        // With kjaq == 0 the macrobody-facet field is not written at all, so the
        // surface type and parameter count sit 4 bytes earlier than with kjaq == 1.
        // Reading `info` unconditionally (as this crate did up to v0.0.6) shifted both
        // of them by one slot.
        let surfaces = [(1, 999, 15, [900.0].as_slice())];
        let bytes = synthetic_mcnp_header(1, 0, 0, &surfaces);
        let (params, _offset) = parse_header(&bytes).unwrap();
        assert_eq!(params.surfaces[0].info, -1);
        assert_eq!(params.surfaces[0].surface_type, 15);
        assert_eq!(params.surfaces[0].parameters, vec![900.0]);
    }

    #[test]
    fn rejects_a_surface_record_whose_length_contradicts_its_parameter_count() {
        let mut bytes = synthetic_mcnp_header(1, 0, 0, &[(1, -1, 15, [900.0].as_slice())]);
        // Overwrite the parameter count of the (single) surface record with an absurd
        // value, the way a misparsed layout would: the length check must reject it
        // rather than trying to reserve gigabytes for the parameter vector.
        let params_offset = bytes
            .windows(8)
            .position(|w| w == 900.0f64.to_le_bytes())
            .expect("the surface record holds its single parameter");
        let count_offset = params_offset - 4;
        bytes[count_offset..count_offset + 4].copy_from_slice(&1_546_188_228i32.to_le_bytes());

        let err = parse_header(&bytes).unwrap_err();
        assert!(matches!(
            err,
            RssaError::UnexpectedSurfaceRecordLength {
                index: 0,
                num_parameters: 1_546_188_228,
                ..
            }
        ));
    }

    #[test]
    fn rejects_unknown_identification_signature() {
        let bytes = write_record(b"not a known signature");
        let err = parse_header(&bytes).unwrap_err();
        assert!(matches!(err, RssaError::UnsupportedFormat(_)));
    }

    #[test]
    fn rejects_niwr_greater_than_zero() {
        let bytes = synthetic_mcnp_header(0, 1, 0, &[]);
        let err = parse_header(&bytes).unwrap_err();
        assert!(matches!(err, RssaError::UnsupportedNiwr(1)));
    }

    #[test]
    fn rejects_non_negative_np1() {
        let mut out = Vec::new();
        out.extend(write_record(b"SF_00001"));
        out.extend(write_record(b"version"));
        let mut dims = Vec::new();
        dims.extend_from_slice(&0i64.to_le_bytes()); // np1 >= 0
        dims.extend_from_slice(&10i64.to_le_bytes());
        dims.extend_from_slice(&11i32.to_le_bytes());
        dims.extend_from_slice(&0i32.to_le_bytes());
        dims.extend_from_slice(&5i64.to_le_bytes());
        out.extend(write_record(&dims));

        let err = parse_header(&out).unwrap_err();
        assert!(matches!(err, RssaError::UnsupportedNp1(0)));
    }

    #[test]
    fn rejects_unexpected_nrcd() {
        let mut out = Vec::new();
        out.extend(write_record(b"SF_00001"));
        out.extend(write_record(b"version"));
        let mut dims = Vec::new();
        dims.extend_from_slice(&(-10i64).to_le_bytes());
        dims.extend_from_slice(&10i64.to_le_bytes());
        dims.extend_from_slice(&7i32.to_le_bytes()); // wrong nrcd
        dims.extend_from_slice(&0i32.to_le_bytes());
        dims.extend_from_slice(&5i64.to_le_bytes());
        out.extend(write_record(&dims));

        let err = parse_header(&out).unwrap_err();
        assert!(matches!(
            err,
            RssaError::UnexpectedRecordLength {
                expected: 11,
                found: 7
            }
        ));
    }

    #[test]
    fn d1suned_variant_exposes_last_dump() {
        let mut id = b"d1suned file".to_vec();
        id.extend_from_slice(&99i32.to_le_bytes());
        let mut out = write_record(&id);

        let mut dims = Vec::new();
        dims.extend_from_slice(&(-10i64).to_le_bytes());
        dims.extend_from_slice(&10i64.to_le_bytes());
        dims.extend_from_slice(&11i32.to_le_bytes());
        dims.extend_from_slice(&0i32.to_le_bytes());
        dims.extend_from_slice(&5i64.to_le_bytes());
        out.extend(write_record(&dims));

        let mut extra = Vec::new();
        extra.extend_from_slice(&0i32.to_le_bytes());
        extra.extend_from_slice(&3i32.to_le_bytes());
        extra.extend_from_slice(&0i32.to_le_bytes());
        out.extend(write_record(&extra));

        out.extend(write_record(b"summary"));

        let (params, _offset) = parse_header(&out).unwrap();
        assert_eq!(params.source_code, SourceCode::D1sUned);
        assert_eq!(params.last_dump, Some(99));
    }
}
