//! Metadata structures for TACO format
//!
//! This module contains structures for storing and manipulating metadata
//! about simulations and atoms in molecular dynamics trajectories.

use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::Value;
use std::collections::HashMap;

/// Metadata about the molecular dynamics simulation
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct SimulationMetadata {
    /// Name of the simulation
    pub name: Option<String>,

    /// Description of the simulation
    pub description: Option<String>,

    /// Temperature in Kelvin
    pub temperature: Option<f64>,

    /// Pressure in bars
    pub pressure: Option<f64>,

    /// Simulation ensemble (e.g., "NVT", "NPT")
    pub ensemble: Option<String>,

    /// Name of the software used for the simulation
    pub software: Option<String>,

    /// Version of the software
    pub software_version: Option<String>,

    /// Force field used for the simulation
    pub force_field: Option<String>,

    /// Integrator used for the simulation
    pub integrator: Option<String>,

    /// Time step in femtoseconds
    pub timestep_fs: Option<f64>,

    /// Additional simulation parameters as key-value pairs
    pub custom_parameters: HashMap<String, String>,
}

impl SimulationMetadata {
    /// Create a new simulation metadata instance
    pub fn new() -> Self {
        Default::default()
    }

    /// Add a custom parameter
    pub fn add_parameter(&mut self, key: &str, value: &str) {
        self.custom_parameters
            .insert(key.to_string(), value.to_string());
    }

    /// Get a custom parameter
    pub fn get_parameter(&self, key: &str) -> Option<&String> {
        self.custom_parameters.get(key)
    }
}

/// Metadata about atoms in the simulation
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct AtomMetadata {
    /// Atom names (e.g., "CA", "N", "O")
    pub names: Option<Vec<String>>,

    /// Atom types (e.g., "C_alpha", "N_amide")
    pub types: Option<Vec<String>>,

    /// Atom elements (e.g., "C", "N", "O")
    pub elements: Option<Vec<String>>,

    /// Atomic numbers (e.g., 1 for H, 6 for C, 8 for O)
    pub atomic_numbers: Option<Vec<u8>>,

    /// Atom masses in atomic mass units (u)
    pub masses: Option<Vec<f32>>,

    /// Atom charges in elementary charge units (e)
    pub charges: Option<Vec<f32>>,

    /// Residue numbers
    pub residue_numbers: Option<Vec<i32>>,

    /// Residue names (e.g., "ALA", "GLY")
    pub residue_names: Option<Vec<String>>,

    /// Chain identifiers (e.g., "A", "B")
    pub chain_ids: Option<Vec<String>>,

    /// Segment identifiers
    pub segment_ids: Option<Vec<String>>,

    /// Additional atom properties as key-value pairs
    /// Keys are property names, values are vectors of property values for each atom
    pub custom_properties: HashMap<String, Vec<String>>,
}

impl AtomMetadata {
    /// Create a new atom metadata instance
    pub fn new() -> Self {
        Default::default()
    }

    /// Add a custom property for atoms
    pub fn add_property(&mut self, name: &str, values: Vec<String>) {
        self.custom_properties.insert(name.to_string(), values);
    }

    /// Get a custom property
    pub fn get_property(&self, name: &str) -> Option<&Vec<String>> {
        self.custom_properties.get(name)
    }

    /// Validate that all property vectors have the expected length
    pub fn validate(&self, expected_atom_count: usize) -> bool {
        // Helper function to validate vector length
        fn check_vector_length<T: Sized>(
            opt: &Option<Vec<T>>,
            name: &str,
            expected_atom_count: usize,
        ) -> bool {
            if let Some(vec) = opt
                && vec.len() != expected_atom_count
            {
                eprintln!(
                    "Metadata error: {} has length {} but expected {}",
                    name,
                    vec.len(),
                    expected_atom_count
                );
                return false;
            }
            true
        }

        check_vector_length(&self.names, "atom names", expected_atom_count)
            && check_vector_length(&self.types, "atom types", expected_atom_count)
            && check_vector_length(&self.elements, "atom elements", expected_atom_count)
            && check_vector_length(&self.atomic_numbers, "atomic numbers", expected_atom_count)
            && check_vector_length(&self.masses, "atom masses", expected_atom_count)
            && check_vector_length(&self.charges, "atom charges", expected_atom_count)
            && check_vector_length(
                &self.residue_numbers,
                "residue numbers",
                expected_atom_count,
            )
            && check_vector_length(&self.residue_names, "residue names", expected_atom_count)
            && check_vector_length(&self.chain_ids, "chain IDs", expected_atom_count)
            && check_vector_length(&self.segment_ids, "segment IDs", expected_atom_count)
            && self.custom_properties.iter().all(|(name, vec)| {
                check_vector_length(
                    &Some(vec.to_vec()),
                    &format!("custom property '{name}'"),
                    expected_atom_count,
                )
            })
    }
}

/// Constraint definitions (ASE-compatible serialization)
#[derive(Debug, Clone, PartialEq)]
pub enum Constraint {
    FixAtoms { indices: Vec<usize> },
    Unknown { name: String, kwargs: Value },
}

#[derive(Debug, Deserialize)]
struct ConstraintRepr {
    name: String,
    #[serde(default)]
    kwargs: Value,
}

#[derive(Debug, Deserialize, Serialize)]
struct FixAtomsKwargs {
    indices: Vec<usize>,
}

impl Serialize for Constraint {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeStruct;

        let mut state = serializer.serialize_struct("Constraint", 2)?;
        match self {
            Constraint::FixAtoms { indices } => {
                state.serialize_field("name", "FixAtoms")?;
                let kwargs = FixAtomsKwargs {
                    indices: indices.clone(),
                };
                state.serialize_field("kwargs", &kwargs)?;
            }
            Constraint::Unknown { name, kwargs } => {
                state.serialize_field("name", name)?;
                state.serialize_field("kwargs", kwargs)?;
            }
        }
        state.end()
    }
}

impl<'de> Deserialize<'de> for Constraint {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let repr = ConstraintRepr::deserialize(deserializer)?;
        match repr.name.as_str() {
            "FixAtoms" => {
                let kwargs: FixAtomsKwargs =
                    serde_json::from_value(repr.kwargs).map_err(serde::de::Error::custom)?;
                Ok(Constraint::FixAtoms {
                    indices: kwargs.indices,
                })
            }
            _ => Ok(Constraint::Unknown {
                name: repr.name,
                kwargs: repr.kwargs,
            }),
        }
    }
}
