//! Header definition and operations for TACO file format
//!
//! The header contains critical metadata about the trajectory including:
//! - Format version
//! - Number of atoms and frames
//! - Compression settings
//! - Simulation parameters
//! - Atom metadata

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{Read, Write};

use crate::compression::{CompressionSettings, PrecisionMode};
use crate::metadata::{AtomMetadata, Constraint, SimulationMetadata};
use crate::{Error, MAGIC, Result, VERSION};

/// Data type for extra datasets
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ExtraDataType {
    F32,
    F64,
}

impl Default for ExtraDataType {
    fn default() -> Self {
        ExtraDataType::F32
    }
}

/// Codec specification for extra datasets
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "codec", rename_all = "lowercase")]
pub enum ExtraCodec {
    Zstd {
        #[serde(default)]
        precision: Option<PrecisionMode>,
    },
    Fpzip {
        #[serde(default)]
        prec: Option<u8>,
    },
}

impl Default for ExtraCodec {
    fn default() -> Self {
        ExtraCodec::Zstd { precision: None }
    }
}

/// Optional grid metadata for 2D/3D datasets
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct GridMetadata {
    pub origin: [f64; 3],
    pub spacing: [f64; 3],
    #[serde(default)]
    pub axis_order: Option<String>,
    #[serde(default)]
    pub units: Option<String>,
}

/// Specification for extra datasets stored in frames
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct ExtraDatasetSpec {
    #[serde(default)]
    pub ndim: Option<u8>,
    #[serde(default)]
    pub shape: Option<Vec<u32>>,
    #[serde(default)]
    pub dtype: ExtraDataType,
    #[serde(default)]
    pub grid: Option<GridMetadata>,
    #[serde(default)]
    pub codec: ExtraCodec,
    #[serde(default)]
    pub full_frame_interval: Option<u32>,
}

/// File format header containing metadata about the trajectory
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Header {
    /// The TACO format version
    pub version: String,

    /// Number of atoms in the trajectory
    pub num_atoms: u32,

    /// Total number of frames in the trajectory
    pub num_frames: u64,

    /// Time step between frames in picoseconds
    pub time_step: f64,

    /// Periodic boundary conditions along [x, y, z]
    pub pbc: [bool; 3],

    /// Interval between full (non-delta) frames
    pub full_frame_interval: u32,

    /// Simulation metadata (temperature, ensemble, etc.)
    pub simulation_metadata: SimulationMetadata,

    /// Atom-specific metadata (masses, names, elements, etc.)
    pub atom_metadata: AtomMetadata,

    /// Constraints (ASE-compatible), apply to all frames
    #[serde(default)]
    pub constraints: Vec<Constraint>,

    /// Compression settings used for the trajectory
    pub compression_settings: CompressionSettings,

    /// Specifications for extra datasets stored in frames
    #[serde(default)]
    pub extra_datasets: HashMap<String, ExtraDatasetSpec>,

    /// Offset to the start of the frame index table (in bytes)
    pub frame_index_offset: u64,

    /// Flags indicating which data components are present
    /// Bit 0: Positions
    /// Bit 1: Velocities
    /// Bit 2: Forces
    /// Bit 3: Box dimensions
    /// Bits 4-31: Reserved for future use
    pub data_flags: u32,
}

impl Header {
    /// Create a new header with default values
    pub fn new(
        num_atoms: u32,
        time_step: f64,
        simulation_metadata: SimulationMetadata,
        atom_metadata: AtomMetadata,
        compression_settings: CompressionSettings,
    ) -> Self {
        Header {
            version: VERSION.to_string(),
            num_atoms,
            num_frames: 0, // Will be updated during writing
            time_step,
            pbc: [true; 3],           // Default: periodic in all directions
            full_frame_interval: 100, // Default: store a full frame every 100 frames
            simulation_metadata,
            atom_metadata,
            constraints: Vec::new(),
            compression_settings,
            extra_datasets: HashMap::new(),
            frame_index_offset: 0, // Will be updated during writing
            data_flags: 0x0F,      // Default: include all data types
        }
    }

    /// Write the header to a writer, optionally padding the JSON payload to a specific length
    pub fn write<W: Write>(&self, writer: &mut W, padded_length: Option<usize>) -> Result<usize> {
        // Write magic number to identify the file format
        writer.write_all(&MAGIC)?;

        // Serialize header to JSON for maximum compatibility and extensibility
        let mut header_json = serde_json::to_vec(self)?;

        if let Some(len) = padded_length {
            if header_json.len() < len {
                header_json.resize(len, b' ');
            } else if header_json.len() > len {
                return Err(Error::InvalidFormat(format!(
                    "Header grew too large during rewrite: {} > {}. Cannot rewrite safely.",
                    header_json.len(),
                    len
                )));
            }
        }

        // Write length of header data as u32
        writer.write_u32::<LittleEndian>(header_json.len() as u32)?;

        // Write the serialized header
        writer.write_all(&header_json)?;

        Ok(header_json.len())
    }

    /// Read the header from a reader
    pub fn read<R: Read>(reader: &mut R) -> Result<(Self, usize)> {
        // Check magic number
        let mut magic = [0u8; 4];
        reader.read_exact(&mut magic)?;

        if magic != MAGIC {
            return Err(Error::InvalidFormat(format!(
                "Invalid file format. Expected TACO magic number, found: {magic:?}"
            )));
        }

        // Read the length of the header data
        let header_length = reader.read_u32::<LittleEndian>()?;

        // Read the serialized header
        let mut header_data = vec![0u8; header_length as usize];
        reader.read_exact(&mut header_data)?;

        // Deserialize the header
        let header: Header = serde_json::from_slice(&header_data)
            .map_err(|e| Error::Deserialization(format!("Failed to deserialize header: {e}")))?;

        // Validate version
        if !is_compatible_version(&header.version) {
            return Err(Error::InvalidFormat(format!(
                "Incompatible format version: {}. This library supports: {}",
                header.version, VERSION
            )));
        }

        Ok((header, header_length as usize))
    }
}

/// Check if the given version is compatible with the current version
fn is_compatible_version(version: &str) -> bool {
    // Accept 0.1.x and 0.2.x versions; in the future, relax to match major.minor
    version.starts_with("0.1.") || version.starts_with("0.2.")
}

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

    #[test]
    fn test_header_roundtrip_without_constraints() {
        let header = Header::new(
            3,
            0.001,
            SimulationMetadata::default(),
            AtomMetadata::default(),
            CompressionSettings::default(),
        );

        let json = serde_json::to_string(&header).expect("Failed to serialize header");
        let decoded: Header = serde_json::from_str(&json).expect("Failed to deserialize header");

        assert!(decoded.constraints.is_empty());
        assert_eq!(decoded.num_atoms, 3);
    }

    #[test]
    fn test_header_deserialize_without_constraints_field() {
        let header = Header::new(
            4,
            0.002,
            SimulationMetadata::default(),
            AtomMetadata::default(),
            CompressionSettings::default(),
        );

        let mut value = serde_json::to_value(&header).expect("Failed to serialize header to value");
        if let serde_json::Value::Object(map) = &mut value {
            map.remove("constraints");
        }

        let decoded: Header =
            serde_json::from_value(value).expect("Failed to deserialize header value");

        assert!(decoded.constraints.is_empty());
        assert_eq!(decoded.num_atoms, 4);
    }

    #[test]
    fn test_header_roundtrip_with_constraints() {
        let mut header = Header::new(
            5,
            0.002,
            SimulationMetadata::default(),
            AtomMetadata::default(),
            CompressionSettings::default(),
        );

        header.constraints = vec![Constraint::FixAtoms {
            indices: vec![0, 2, 4],
        }];

        let json = serde_json::to_string(&header).expect("Failed to serialize header");
        let decoded: Header = serde_json::from_str(&json).expect("Failed to deserialize header");

        assert_eq!(decoded.constraints.len(), 1);
        match &decoded.constraints[0] {
            Constraint::FixAtoms { indices } => {
                assert_eq!(indices, &vec![0, 2, 4]);
            }
            _ => panic!("Expected FixAtoms constraint"),
        }
    }
}
