use std::collections::BTreeMap;

use super::*;
use crate::{
    parse::{Parse, ParseError, RawParseError},
    v1,
};
use anyhow::Result;

impl Parse for v1::Equality {
    type Output = Equality;
    type Context = ();
    fn parse(self, _: &Self::Context) -> Result<Self::Output, ParseError> {
        match self {
            v1::Equality::EqualToZero => Ok(Equality::EqualToZero),
            v1::Equality::LessThanOrEqualToZero => Ok(Equality::LessThanOrEqualToZero),
            _ => Err(RawParseError::UnknownEnumValue {
                enum_name: "ommx.v1.Equality",
                value: self as i32,
            }
            .into()),
        }
    }
}

impl From<Equality> for v1::Equality {
    fn from(value: Equality) -> Self {
        match value {
            Equality::EqualToZero => v1::Equality::EqualToZero,
            Equality::LessThanOrEqualToZero => v1::Equality::LessThanOrEqualToZero,
        }
    }
}

impl From<Equality> for i32 {
    fn from(equality: Equality) -> Self {
        v1::Equality::from(equality).into()
    }
}

impl Parse for v1::Constraint {
    type Output = (ConstraintID, Constraint<Created>, ConstraintMetadata);
    type Context = ();

    fn parse(self, _: &Self::Context) -> Result<Self::Output, ParseError> {
        let message = "ommx.v1.Constraint";
        let id = ConstraintID(self.id);
        let equality = self.equality().parse_as(&(), message, "equality")?;
        let metadata = ConstraintMetadata {
            name: self.name,
            subscripts: self.subscripts,
            parameters: self.parameters.into_iter().collect(),
            description: self.description,
            provenance: Vec::new(),
        };
        let function = self
            .function
            .ok_or(RawParseError::MissingField {
                message,
                field: "function",
            })?
            .parse_as(&(), message, "function")?;
        Ok((
            id,
            Constraint {
                equality,
                stage: CreatedData { function },
            },
            metadata,
        ))
    }
}

impl Parse for v1::RemovedConstraint {
    type Output = (
        ConstraintID,
        Constraint<Created>,
        ConstraintMetadata,
        RemovedReason,
    );
    type Context = ();

    fn parse(self, _: &Self::Context) -> Result<Self::Output, ParseError> {
        let message = "ommx.v1.RemovedConstraint";
        let (id, constraint, metadata) = self
            .constraint
            .ok_or(RawParseError::MissingField {
                message,
                field: "constraint",
            })?
            .parse_as(&(), message, "constraint")?;
        let removed_reason = RemovedReason {
            reason: self.removed_reason,
            parameters: self.removed_reason_parameters.into_iter().collect(),
        };
        Ok((id, constraint, metadata, removed_reason))
    }
}

impl Parse for Vec<v1::Constraint> {
    type Output = (
        BTreeMap<ConstraintID, Constraint<Created>>,
        ConstraintMetadataStore<ConstraintID>,
    );
    type Context = ();
    fn parse(self, _: &Self::Context) -> Result<Self::Output, ParseError> {
        let mut constraints = BTreeMap::default();
        let mut metadata_store = ConstraintMetadataStore::default();
        for c in self {
            let (id, c, metadata): (ConstraintID, Constraint<Created>, ConstraintMetadata) =
                c.parse(&())?;
            if constraints.insert(id, c).is_some() {
                return Err(RawParseError::InvalidInstance(format!(
                    "Duplicated constraint ID is found in definition: {id:?}"
                ))
                .into());
            }
            metadata_store.insert(id, metadata);
        }
        Ok((constraints, metadata_store))
    }
}

impl Parse for Vec<v1::RemovedConstraint> {
    type Output = BTreeMap<ConstraintID, (Constraint<Created>, ConstraintMetadata, RemovedReason)>;
    type Context = BTreeMap<ConstraintID, Constraint<Created>>;
    fn parse(self, constraints: &Self::Context) -> Result<Self::Output, ParseError> {
        let mut removed_constraints = BTreeMap::default();
        for c in self {
            let (id, constraint, metadata, reason) = c.parse(&())?;
            if constraints.contains_key(&id) {
                return Err(RawParseError::InvalidInstance(format!(
                    "Duplicated constraint ID is found in definition: {id:?}"
                ))
                .into());
            }
            if removed_constraints
                .insert(id, (constraint, metadata, reason))
                .is_some()
            {
                return Err(RawParseError::InvalidInstance(format!(
                    "Duplicated constraint ID is found in definition: {id:?}"
                ))
                .into());
            }
        }
        Ok(removed_constraints)
    }
}

impl Parse for v1::EvaluatedConstraint {
    type Output = (
        ConstraintID,
        EvaluatedConstraint,
        ConstraintMetadata,
        Option<RemovedReason>,
    );
    type Context = ();

    fn parse(self, _: &Self::Context) -> Result<Self::Output, ParseError> {
        let message = "ommx.v1.EvaluatedConstraint";

        let equality = self.equality().parse_as(&(), message, "equality")?;

        let metadata = ConstraintMetadata {
            name: self.name,
            subscripts: self.subscripts,
            parameters: self.parameters.into_iter().collect(),
            description: self.description,
            provenance: Vec::new(),
        };

        let feasible = match equality {
            Equality::EqualToZero => self.evaluated_value.abs() < *crate::ATol::default(),
            Equality::LessThanOrEqualToZero => self.evaluated_value < *crate::ATol::default(),
        };

        let removed_reason = self.removed_reason.map(|reason| RemovedReason {
            reason,
            parameters: self.removed_reason_parameters.into_iter().collect(),
        });

        Ok((
            ConstraintID(self.id),
            Constraint {
                equality,
                stage: EvaluatedData {
                    evaluated_value: self.evaluated_value,
                    dual_variable: self.dual_variable,
                    feasible,
                    used_decision_variable_ids: self
                        .used_decision_variable_ids
                        .into_iter()
                        .map(VariableID::from)
                        .collect(),
                },
            },
            metadata,
            removed_reason,
        ))
    }
}

impl Parse for v1::SampledConstraint {
    type Output = (
        ConstraintID,
        SampledConstraint,
        ConstraintMetadata,
        Option<RemovedReason>,
    );
    type Context = ();

    fn parse(self, _: &Self::Context) -> Result<Self::Output, ParseError> {
        let message = "ommx.v1.SampledConstraint";

        let equality = self.equality().parse_as(&(), message, "equality")?;

        // Parse evaluated_values
        let evaluated_values = self
            .evaluated_values
            .ok_or(RawParseError::MissingField {
                message,
                field: "evaluated_values",
            })?
            .parse_as(&(), message, "evaluated_values")?;

        let metadata = ConstraintMetadata {
            name: self.name,
            subscripts: self.subscripts,
            parameters: self.parameters.into_iter().collect(),
            description: self.description,
            provenance: Vec::new(),
        };

        let removed_reason = self.removed_reason.map(|reason| RemovedReason {
            reason,
            parameters: self.removed_reason_parameters.into_iter().collect(),
        });

        Ok((
            ConstraintID(self.id),
            Constraint {
                equality,
                stage: SampledData {
                    evaluated_values,
                    dual_variables: None,
                    feasible: self
                        .feasible
                        .into_iter()
                        .map(|(id, value)| (SampleID::from(id), value))
                        .collect(),
                    used_decision_variable_ids: self
                        .used_decision_variable_ids
                        .into_iter()
                        .map(VariableID::from)
                        .collect(),
                },
            },
            metadata,
            removed_reason,
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::v1;
    use maplit::btreeset;

    #[test]
    fn error_message() {
        let out: Result<
            (
                ConstraintID,
                Constraint<Created>,
                ConstraintMetadata,
                RemovedReason,
            ),
            ParseError,
        > = v1::RemovedConstraint {
            constraint: Some(v1::Constraint {
                id: 1,
                function: Some(v1::Function { function: None }),
                equality: v1::Equality::EqualToZero as i32,
                ..Default::default()
            }),
            removed_reason: "reason".to_string(),
            removed_reason_parameters: Default::default(),
        }
        .parse(&());

        insta::assert_snapshot!(out.unwrap_err(), @r###"
        Traceback for OMMX Message parse error:
        └─ommx.v1.RemovedConstraint[constraint]
          └─ommx.v1.Constraint[function]
        Unsupported ommx.v1.Function is found. It is created by a newer version of OMMX SDK.
        "###);
    }

    #[test]
    fn test_evaluated_constraint_parse() {
        let v1_constraint = v1::EvaluatedConstraint {
            id: 42,
            equality: v1::Equality::EqualToZero as i32,
            evaluated_value: 1.5,
            used_decision_variable_ids: vec![1, 2, 3],
            subscripts: vec![10, 20],
            parameters: [("key1".to_string(), "value1".to_string())]
                .iter()
                .cloned()
                .collect(),
            name: Some("test_constraint".to_string()),
            description: Some("A test constraint".to_string()),
            dual_variable: Some(0.5),
            removed_reason: None,
            removed_reason_parameters: Default::default(),
        };

        let (id, parsed, metadata, removed_reason): (
            ConstraintID,
            EvaluatedConstraint,
            ConstraintMetadata,
            _,
        ) = v1_constraint.parse(&()).unwrap();
        assert!(removed_reason.is_none());

        assert_eq!(id, ConstraintID(42));
        assert_eq!(parsed.equality, Equality::EqualToZero);
        assert_eq!(parsed.stage.evaluated_value, 1.5);
        assert_eq!(parsed.stage.dual_variable, Some(0.5));
        assert_eq!(
            parsed.stage.used_decision_variable_ids,
            btreeset! {1.into(), 2.into(), 3.into()}
        );
        // Metadata is now returned as a separate value alongside the parsed
        // constraint (drained into the SoA store at the collection level).
        assert_eq!(metadata.name, Some("test_constraint".to_string()));
        assert_eq!(metadata.description, Some("A test constraint".to_string()));
        assert_eq!(metadata.subscripts, vec![10, 20]);
        // feasible should be false because 1.5 > ATol::default() for EqualToZero constraint
        assert!(!parsed.stage.feasible);
    }
}
