use crate::dimension::Dimension;
use crate::dtype::DType;
use crate::error::{ArrDbError, Result};

/// Whether coordinate values are monotonically ordered.
///
/// Monotonic coordinates enable binary search for label→index lookup.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Monotonicity {
    Ascending,
    Descending,
    Unknown,
}

/// Metadata for a coordinate variable attached to a dimension.
///
/// Coordinate variables provide real-valued labels (e.g., latitude in degrees,
/// time as timestamps) for a dimension's integer indices. The actual coordinate
/// values are stored in a [`CoordinateIndex`](crate::coord_index::CoordinateIndex).
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CoordinateVar {
    /// Must match a dimension name in the parent `DatasetSchema`.
    pub dim_name: String,
    /// Type of coordinate values (Float64 for lat/lon, Timestamp for time, etc.).
    pub dtype: DType,
    /// Whether values are monotonically ordered (enables binary search).
    pub monotonicity: Monotonicity,
    /// Number of coordinate values (must equal the dimension extent).
    pub len: u64,
}

/// A data variable within a dataset — references dimensions by index.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct DataVar {
    /// Variable name (e.g., "temperature", "2m_temperature").
    pub name: String,
    /// Cell data type.
    pub dtype: DType,
    /// Indices into `DatasetSchema.dimensions` — the dimensions this variable spans.
    pub dim_indices: Vec<usize>,
}

/// Schema for a dataset — a collection of related multi-dimensional variables
/// sharing named dimensions and coordinate variables.
///
/// This is the xarray-equivalent model: a group of N-dimensional data variables
/// that share named dimensions, with coordinate variables providing real-valued
/// labels for those dimensions.
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct DatasetSchema {
    /// Dataset name (unique within a dataset registry).
    pub name: String,
    /// Union of all dimensions across all data variables.
    pub dimensions: Vec<Dimension>,
    /// Data variables (the arrays with actual data).
    pub data_vars: Vec<DataVar>,
    /// Coordinate variable metadata (one per dimension that has coordinates).
    pub coords: Vec<CoordinateVar>,
}

impl DatasetSchema {
    /// Creates a new dataset schema, validating constraints.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No dimensions are provided
    /// - No data variables are provided
    /// - A data variable references an out-of-range dimension index
    /// - A coordinate variable's `dim_name` doesn't match any dimension
    /// - Duplicate data variable names
    /// - Duplicate coordinate variable dim_names
    pub fn new(
        name: impl Into<String>,
        dimensions: Vec<Dimension>,
        data_vars: Vec<DataVar>,
        coords: Vec<CoordinateVar>,
    ) -> Result<Self> {
        if dimensions.is_empty() {
            return Err(ArrDbError::Schema(
                "dataset must have at least one dimension".into(),
            ));
        }
        if data_vars.is_empty() {
            return Err(ArrDbError::Schema(
                "dataset must have at least one data variable".into(),
            ));
        }

        // Validate data variable dimension indices.
        for var in &data_vars {
            for &idx in &var.dim_indices {
                if idx >= dimensions.len() {
                    return Err(ArrDbError::Schema(format!(
                        "data variable '{}' references dimension index {} but only {} dimensions exist",
                        var.name,
                        idx,
                        dimensions.len()
                    )));
                }
            }
        }

        // Check for duplicate data variable names.
        let mut var_names = std::collections::HashSet::new();
        for var in &data_vars {
            if !var_names.insert(&var.name) {
                return Err(ArrDbError::Schema(format!(
                    "duplicate data variable name: '{}'",
                    var.name
                )));
            }
        }

        // Validate coordinate variables.
        let mut coord_dims = std::collections::HashSet::new();
        for coord in &coords {
            // Must reference an existing dimension.
            let dim = dimensions.iter().find(|d| d.name == coord.dim_name);
            match dim {
                None => {
                    return Err(ArrDbError::Schema(format!(
                        "coordinate '{}' references unknown dimension",
                        coord.dim_name
                    )));
                }
                Some(d) => {
                    if coord.len != d.extent() {
                        return Err(ArrDbError::Schema(format!(
                            "coordinate '{}' has len {} but dimension extent is {}",
                            coord.dim_name,
                            coord.len,
                            d.extent()
                        )));
                    }
                }
            }
            if !coord_dims.insert(&coord.dim_name) {
                return Err(ArrDbError::Schema(format!(
                    "duplicate coordinate for dimension: '{}'",
                    coord.dim_name
                )));
            }
        }

        Ok(Self {
            name: name.into(),
            dimensions,
            data_vars,
            coords,
        })
    }

    /// Number of dimensions in this dataset.
    pub fn ndim(&self) -> usize {
        self.dimensions.len()
    }

    /// Look up a dimension by name.
    pub fn dimension(&self, name: &str) -> Option<&Dimension> {
        self.dimensions.iter().find(|d| d.name == name)
    }

    /// Look up a dimension index by name.
    pub fn dimension_index(&self, name: &str) -> Option<usize> {
        self.dimensions.iter().position(|d| d.name == name)
    }

    /// Look up a data variable by name.
    pub fn data_var(&self, name: &str) -> Option<&DataVar> {
        self.data_vars.iter().find(|v| v.name == name)
    }

    /// Look up a coordinate variable by dimension name.
    pub fn coord(&self, dim_name: &str) -> Option<&CoordinateVar> {
        self.coords.iter().find(|c| c.dim_name == dim_name)
    }

    /// Returns the dimension names for a given data variable.
    pub fn var_dim_names(&self, var: &DataVar) -> Vec<&str> {
        var.dim_indices
            .iter()
            .map(|&i| self.dimensions[i].name.as_str())
            .collect()
    }
}

impl std::fmt::Display for DatasetSchema {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "DATASET {} (", self.name)?;
        writeln!(f, "  Dimensions:")?;
        for dim in &self.dimensions {
            writeln!(f, "    {dim}")?;
        }
        writeln!(f, "  Coordinates:")?;
        for coord in &self.coords {
            let mono = match coord.monotonicity {
                Monotonicity::Ascending => " [ascending]",
                Monotonicity::Descending => " [descending]",
                Monotonicity::Unknown => "",
            };
            writeln!(
                f,
                "    {} {} (len={}){mono}",
                coord.dim_name, coord.dtype, coord.len
            )?;
        }
        writeln!(f, "  Data variables:")?;
        for var in &self.data_vars {
            let dims: Vec<&str> = self.var_dim_names(var);
            writeln!(f, "    {} {} [{}]", var.name, var.dtype, dims.join(", "))?;
        }
        write!(f, ")")
    }
}

impl std::fmt::Display for Monotonicity {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Monotonicity::Ascending => write!(f, "ascending"),
            Monotonicity::Descending => write!(f, "descending"),
            Monotonicity::Unknown => write!(f, "unknown"),
        }
    }
}

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

    fn era5_schema() -> DatasetSchema {
        DatasetSchema::new(
            "era5",
            vec![
                Dimension::new("time", DType::Timestamp(TimeUnit::Second), 0, 8759, 720),
                Dimension::new("latitude", DType::Float64, 0, 720, 180),
                Dimension::new("longitude", DType::Float64, 0, 1439, 360),
            ],
            vec![
                DataVar {
                    name: "temperature".into(),
                    dtype: DType::Float32,
                    dim_indices: vec![0, 1, 2],
                },
                DataVar {
                    name: "humidity".into(),
                    dtype: DType::Float32,
                    dim_indices: vec![0, 1, 2],
                },
            ],
            vec![
                CoordinateVar {
                    dim_name: "time".into(),
                    dtype: DType::Timestamp(TimeUnit::Second),
                    monotonicity: Monotonicity::Ascending,
                    len: 8760,
                },
                CoordinateVar {
                    dim_name: "latitude".into(),
                    dtype: DType::Float64,
                    monotonicity: Monotonicity::Descending,
                    len: 721,
                },
                CoordinateVar {
                    dim_name: "longitude".into(),
                    dtype: DType::Float64,
                    monotonicity: Monotonicity::Ascending,
                    len: 1440,
                },
            ],
        )
        .unwrap()
    }

    #[test]
    fn dataset_schema_basic() {
        let ds = era5_schema();
        assert_eq!(ds.ndim(), 3);
        assert_eq!(ds.data_vars.len(), 2);
        assert_eq!(ds.coords.len(), 3);
    }

    #[test]
    fn dataset_schema_lookups() {
        let ds = era5_schema();
        assert!(ds.dimension("time").is_some());
        assert!(ds.dimension("missing").is_none());
        assert_eq!(ds.dimension_index("latitude"), Some(1));
        assert!(ds.data_var("temperature").is_some());
        assert!(ds.data_var("pressure").is_none());
        assert!(ds.coord("time").is_some());
        assert!(ds.coord("missing").is_none());
    }

    #[test]
    fn dataset_var_dim_names() {
        let ds = era5_schema();
        let temp = ds.data_var("temperature").unwrap();
        let names = ds.var_dim_names(temp);
        assert_eq!(names, vec!["time", "latitude", "longitude"]);
    }

    #[test]
    fn dataset_no_dims() {
        let r = DatasetSchema::new("bad", vec![], vec![], vec![]);
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_no_vars() {
        let r = DatasetSchema::new(
            "bad",
            vec![Dimension::new("x", DType::Float64, 0, 9, 10)],
            vec![],
            vec![],
        );
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_bad_dim_index() {
        let r = DatasetSchema::new(
            "bad",
            vec![Dimension::new("x", DType::Float64, 0, 9, 10)],
            vec![DataVar {
                name: "v".into(),
                dtype: DType::Float32,
                dim_indices: vec![0, 5], // index 5 doesn't exist
            }],
            vec![],
        );
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_bad_coord_dim() {
        let r = DatasetSchema::new(
            "bad",
            vec![Dimension::new("x", DType::Float64, 0, 9, 10)],
            vec![DataVar {
                name: "v".into(),
                dtype: DType::Float32,
                dim_indices: vec![0],
            }],
            vec![CoordinateVar {
                dim_name: "nonexistent".into(),
                dtype: DType::Float64,
                monotonicity: Monotonicity::Unknown,
                len: 10,
            }],
        );
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_coord_len_mismatch() {
        let r = DatasetSchema::new(
            "bad",
            vec![Dimension::new("x", DType::Float64, 0, 9, 10)],
            vec![DataVar {
                name: "v".into(),
                dtype: DType::Float32,
                dim_indices: vec![0],
            }],
            vec![CoordinateVar {
                dim_name: "x".into(),
                dtype: DType::Float64,
                monotonicity: Monotonicity::Ascending,
                len: 99, // doesn't match extent=10
            }],
        );
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_duplicate_var_names() {
        let r = DatasetSchema::new(
            "bad",
            vec![Dimension::new("x", DType::Float64, 0, 9, 10)],
            vec![
                DataVar {
                    name: "v".into(),
                    dtype: DType::Float32,
                    dim_indices: vec![0],
                },
                DataVar {
                    name: "v".into(),
                    dtype: DType::Float64,
                    dim_indices: vec![0],
                },
            ],
            vec![],
        );
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_duplicate_coord_dims() {
        let r = DatasetSchema::new(
            "bad",
            vec![Dimension::new("x", DType::Float64, 0, 9, 10)],
            vec![DataVar {
                name: "v".into(),
                dtype: DType::Float32,
                dim_indices: vec![0],
            }],
            vec![
                CoordinateVar {
                    dim_name: "x".into(),
                    dtype: DType::Float64,
                    monotonicity: Monotonicity::Ascending,
                    len: 10,
                },
                CoordinateVar {
                    dim_name: "x".into(),
                    dtype: DType::Float64,
                    monotonicity: Monotonicity::Descending,
                    len: 10,
                },
            ],
        );
        assert!(matches!(r, Err(ArrDbError::Schema(_))));
    }

    #[test]
    fn dataset_display() {
        let ds = era5_schema();
        let text = ds.to_string();
        assert!(text.contains("DATASET era5"));
        assert!(text.contains("temperature"));
        assert!(text.contains("latitude"));
        assert!(text.contains("[ascending]"));
        assert!(text.contains("[descending]"));
    }

    #[test]
    fn monotonicity_display() {
        assert_eq!(Monotonicity::Ascending.to_string(), "ascending");
        assert_eq!(Monotonicity::Descending.to_string(), "descending");
        assert_eq!(Monotonicity::Unknown.to_string(), "unknown");
    }

    #[cfg(feature = "serde")]
    #[test]
    fn dataset_schema_serde_roundtrip() {
        let ds = era5_schema();
        let json = serde_json::to_string(&ds).unwrap();
        let back: DatasetSchema = serde_json::from_str(&json).unwrap();
        assert_eq!(back.name, ds.name);
        assert_eq!(back.dimensions.len(), ds.dimensions.len());
        assert_eq!(back.data_vars.len(), ds.data_vars.len());
        assert_eq!(back.coords.len(), ds.coords.len());
    }
}
