use anyhow::anyhow;
use anyhow::Result;
use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use uuid::Uuid;

use crate::arithmetic_operator::ArithmeticOperator;
use crate::comparison_operator::ComparisonOperator;
use crate::node_props::NodePropsType;
use crate::split::DbAssignmentEntry;

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(tag = "type")]
pub enum Expression {
    #[serde(rename = "ApplicationExpression")]
    Application(ApplicationExpression),

    #[serde(rename = "ObjectExpression")]
    Object(ObjectExpression),

    #[serde(rename = "BooleanExpression")]
    Boolean(BooleanExpression),

    #[serde(rename = "FunctionExpression")]
    Function(FunctionExpression),

    #[serde(rename = "SwitchExpression")]
    Switch(SwitchExpression),

    #[serde(rename = "ComparisonExpression")]
    Comparison(ComparisonExpression),

    #[serde(rename = "GetFieldExpression")]
    GetField(GetFieldExpression),

    #[serde(rename = "ListExpression")]
    List(ListExpression),

    #[serde(rename = "IntExpression")]
    Int(IntExpression),

    #[serde(rename = "FloatExpression")]
    Float(FloatExpression),

    #[serde(rename = "StringExpression")]
    String(StringExpression),

    #[serde(rename = "VariableExpression")]
    Variable(VariableExpression),

    #[serde(rename = "LogEventExpression")]
    LogEvent(LogEventExpression),

    #[serde(rename = "SplitExpression")]
    Split(SplitExpression),

    #[serde(rename = "EnumSwitchExpression")]
    EnumSwitch(EnumSwitchExpression),

    #[serde(rename = "ArithmeticExpression")]
    Arithmetic(ArithmeticExpression),

    #[serde(rename = "NoOpExpression")]
    NoOp(NoOpExpression),

    #[serde(rename = "EnumExpression")]
    Enum(EnumExpression),
}

impl Expression {
    pub fn merge_logs(mut self, logs: &ReductionLogs) -> Self {
        self.set_logs(ReductionLogs::merge(self.get_logs(), logs));
        self
    }

    pub fn get_logs(&self) -> &ReductionLogs {
        match self {
            Expression::Application(e) => &e.logs,
            Expression::Object(e) => &e.logs,
            Expression::Function(e) => &e.logs,
            Expression::Switch(e) => &e.logs,
            Expression::GetField(e) => &e.logs,
            Expression::List(e) => &e.logs,
            Expression::Int(e) => &e.logs,
            Expression::Float(e) => &e.logs,
            Expression::Comparison(e) => &e.logs,
            Expression::String(e) => &e.logs,
            Expression::Variable(e) => &e.logs,
            Expression::Boolean(e) => &e.logs,
            Expression::NoOp(e) => &e.logs,
            Expression::LogEvent(e) => &e.logs,
            Expression::Split(e) => &e.logs,
            Expression::EnumSwitch(e) => &e.logs,
            Expression::Arithmetic(e) => &e.logs,
            Expression::Enum(e) => &e.logs,
        }
    }

    fn set_logs(&mut self, logs: ReductionLogs) {
        match self {
            Expression::Application(e) => e.logs = logs,
            Expression::Object(e) => e.logs = logs,
            Expression::Function(e) => e.logs = logs,
            Expression::Switch(e) => e.logs = logs,
            Expression::GetField(e) => e.logs = logs,
            Expression::List(e) => e.logs = logs,
            Expression::Int(e) => e.logs = logs,
            Expression::Float(e) => e.logs = logs,
            Expression::Comparison(e) => e.logs = logs,
            Expression::String(e) => e.logs = logs,
            Expression::Variable(e) => e.logs = logs,
            Expression::Boolean(e) => e.logs = logs,
            Expression::NoOp(e) => e.logs = logs,
            Expression::LogEvent(e) => e.logs = logs,
            Expression::Split(e) => e.logs = logs,
            Expression::EnumSwitch(e) => e.logs = logs,
            Expression::Arithmetic(e) => e.logs = logs,
            Expression::Enum(e) => e.logs = logs,
        }
    }

    pub fn is_transient(&self) -> bool {
        match self {
            Expression::Application(_e) => false,
            Expression::Object(e) => e.is_transient,
            Expression::Function(_e) => false,
            Expression::Switch(_e) => false,
            Expression::GetField(_e) => false,
            Expression::List(e) => e.is_transient,
            Expression::Int(e) => e.is_transient,
            Expression::Float(e) => e.is_transient,
            Expression::Comparison(_e) => false,
            Expression::String(e) => e.is_transient,
            Expression::Variable(_e) => false,
            Expression::Boolean(e) => e.is_transient,
            Expression::NoOp(e) => e.is_transient,
            Expression::LogEvent(_e) => false,
            Expression::Split(_e) => false,
            Expression::EnumSwitch(_e) => false,
            Expression::Arithmetic(_e) => false,
            Expression::Enum(_e) => false,
        }
    }

    pub fn get_id(&self) -> &str {
        match self {
            Expression::Application(e) => &e.id,
            Expression::Object(e) => &e.id,
            Expression::Function(e) => &e.id,
            Expression::Switch(e) => &e.id,
            Expression::GetField(e) => &e.id,
            Expression::List(e) => &e.id,
            Expression::Int(e) => &e.id,
            Expression::Float(e) => &e.id,
            Expression::Comparison(e) => &e.id,
            Expression::String(e) => &e.id,
            Expression::Variable(e) => &e.id,
            Expression::Boolean(e) => &e.id,
            Expression::NoOp(e) => &e.id,
            Expression::LogEvent(e) => &e.id,
            Expression::Split(e) => &e.id,
            Expression::EnumSwitch(e) => &e.id,
            Expression::Arithmetic(e) => &e.id,
            Expression::Enum(e) => &e.id,
        }
    }

    pub fn set_id(self, id: String) -> Self {
        match self {
            Expression::Application(mut e) => {
                e.id = id;
                Expression::Application(e)
            }
            Expression::Object(mut e) => {
                e.id = id;
                Expression::Object(e)
            }
            Expression::Function(mut e) => {
                e.id = id;
                Expression::Function(e)
            }
            Expression::Switch(mut e) => {
                e.id = id;
                Expression::Switch(e)
            }
            Expression::GetField(mut e) => {
                e.id = id;
                Expression::GetField(e)
            }
            Expression::List(mut e) => {
                e.id = id;
                Expression::List(e)
            }
            Expression::Int(mut e) => {
                e.id = id;
                Expression::Int(e)
            }
            Expression::Float(mut e) => {
                e.id = id;
                Expression::Float(e)
            }
            Expression::Comparison(mut e) => {
                e.id = id;
                Expression::Comparison(e)
            }
            Expression::String(mut e) => {
                e.id = id;
                Expression::String(e)
            }
            Expression::Variable(mut e) => {
                e.id = id;
                Expression::Variable(e)
            }
            Expression::Boolean(mut e) => {
                e.id = id;
                Expression::Boolean(e)
            }
            Expression::NoOp(mut e) => {
                e.id = id;
                Expression::NoOp(e)
            }
            Expression::LogEvent(mut e) => {
                e.id = id;
                Expression::LogEvent(e)
            }
            Expression::Split(mut e) => {
                e.id = id;
                Expression::Split(e)
            }
            Expression::EnumSwitch(mut e) => {
                e.id = id;
                Expression::EnumSwitch(e)
            }
            Expression::Arithmetic(mut e) => {
                e.id = id;
                Expression::Arithmetic(e)
            }
            Expression::Enum(mut e) => {
                e.id = id;
                Expression::Enum(e)
            }
        }
    }

    pub fn get_node_props_type(&self) -> NodePropsType {
        match self {
            Expression::String(_) => NodePropsType::String,
            Expression::Boolean(_) => NodePropsType::Bool,
            Expression::Object(ObjectExpression {
                object_type_name, ..
            }) => NodePropsType::Object {
                object_type_name: object_type_name.to_string(),
            },
            Expression::Int(_) => NodePropsType::Int,
            Expression::Float(_) => NodePropsType::Float,
            Expression::NoOp(_) => NodePropsType::Void,
            Expression::Enum(_) => NodePropsType::Enum,
            _ => NodePropsType::NotFullyReduced,
        }
    }
}

// TODO: Create a macro to add all expression base fields
#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct ApplicationExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub function: Box<Expression>,
    pub arguments: Vec<Expression>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct ObjectExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub object_type_name: String,
    pub fields: HashMap<String, Expression>,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl ObjectExpression {
    pub fn transient(fields: HashMap<String, Expression>) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            object_type_name: "".to_string(),
            fields,
            logs: ReductionLogs::empty(),
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct BooleanExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub value: bool,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl BooleanExpression {
    pub fn transient(value: bool) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            value,
            logs: ReductionLogs::empty(),
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
pub struct FunctionExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub body: Box<Expression>,
    pub parameters: Vec<FunctionParameter>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
pub struct SwitchExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub control: Box<Expression>,
    pub cases: Vec<SwitchCase>,
    pub default: Box<Expression>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
pub struct ComparisonExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub a: Box<Expression>,
    pub b: Box<Expression>,
    pub operator: ComparisonOperator,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct GetFieldExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub field_path: String,
    pub object: Box<Expression>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct ListExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub items: Vec<Expression>,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl ListExpression {
    pub fn transient(items: Vec<Expression>) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            items,
            logs: ReductionLogs::empty(),
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct IntExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub value: i64,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl IntExpression {
    pub fn transient(value: i64) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            logs: ReductionLogs::empty(),
            value,
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct FloatExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub value: OrderedFloat<f64>,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl FloatExpression {
    pub fn transient(value: f64) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            logs: ReductionLogs::empty(),
            value: value.into(),
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct StringExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub value: String,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl StringExpression {
    pub fn transient(value: String) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            value,
            logs: ReductionLogs::empty(),
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct NoOpExpression {
    pub id: String,
    pub logs: ReductionLogs,

    #[serde(default, skip_serializing)]
    pub is_transient: bool,
}

impl NoOpExpression {
    pub fn transient() -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            logs: ReductionLogs::empty(),
            is_transient: true,
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct EnumExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub value: String,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct VariableExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub variable_id: String,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct LogEventExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub event_type_id: String,
    pub unit_id: Box<Expression>,
    pub features_mapping: HashMap<String, Expression>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct SplitExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub split_id: String,
    pub dimension_id: String,
    pub expose: Box<Expression>,
    pub unit_id: Box<Expression>,
    pub dimension_mapping: DimensionMapping,
    pub features_mapping: HashMap<String, Expression>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct EnumSwitchExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub control: Box<Expression>,
    pub cases: HashMap<String, Expression>,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct ArithmeticExpression {
    pub id: String,
    pub logs: ReductionLogs,

    pub a: Box<Expression>,
    pub b: Box<Expression>,
    pub operator: ArithmeticOperator,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum DimensionMapping {
    Discrete { cases: HashMap<String, Expression> },
    Continuous { function: Box<Expression> },
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
pub struct SwitchCase {
    pub when: Expression,
    pub then: Expression,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
pub struct FunctionParameter {
    pub id: String,
    pub name: String,
}

impl TryFrom<Value> for Expression {
    type Error = anyhow::Error;

    /// Construct an anonymous expression from a JSON value
    fn try_from(value: Value) -> Result<Self> {
        match value {
            Value::Bool(value) => Ok(Expression::Boolean(BooleanExpression::transient(value))),
            Value::String(value) => Ok(Expression::String(StringExpression::transient(value))),
            Value::Number(value) => match value.is_i64() {
                true => Ok(Expression::Int(IntExpression::transient(
                    value
                        .as_i64()
                        .ok_or(anyhow!("Failed to fit integer inside i64"))?,
                ))),
                false => Ok(Expression::Float(FloatExpression::transient(
                    value
                        .as_f64()
                        .ok_or(anyhow!("Failed to represent number as f64"))?,
                ))),
            },
            Value::Array(array) => {
                let items: Vec<Expression> = array
                    .into_iter()
                    .map(|value| value.try_into())
                    .collect::<Result<_>>()?;
                Ok(Expression::List(ListExpression::transient(items)))
            }
            Value::Object(object) => {
                let fields: HashMap<String, Expression> = object
                    .into_iter()
                    .map(|(key, value)| match value.try_into() {
                        Ok(value) => Ok((key, value)),
                        Err(error) => Err(error),
                    })
                    .collect::<Result<_>>()?;
                Ok(Expression::Object(ObjectExpression::transient(fields)))
            }
            Value::Null => Err(anyhow!(
                "Cannot convert a null JSON value into an expression"
            )),
        }
    }
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
pub struct ReductionLogs {
    pub evaluations: HashMap<String, u64>,
    pub events: HashMap<String, u64>,
    pub exposures: HashMap<String, u64>,
}

fn sort_alphabetically<T: Serialize, S: serde::Serializer>(
    value: &T,
    serializer: S,
) -> Result<S::Ok, S::Error> {
    let value = serde_json::to_value(value).map_err(serde::ser::Error::custom)?;
    value.serialize(serializer)
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
// N.B. Keep the keys in sorted order so the serialization is 'stable'
pub struct Exposure {
    #[serde(serialize_with = "sort_alphabetically")]
    pub assignment: HashMap<String, DbAssignmentEntry>,

    #[serde(serialize_with = "sort_alphabetically")]
    pub feature_values: HashMap<String, Value>,

    pub split_id: String,
    pub unit_id: String,
}

#[derive(Deserialize, Serialize, Eq, PartialEq, Debug, Clone)]
#[serde(rename_all = "camelCase")]
// N.B. Keep the keys in sorted order so the serialization is 'stable'
pub struct Event {
    pub event_type_id: String,

    #[serde(serialize_with = "sort_alphabetically")]
    pub feature_values: HashMap<String, Value>,

    pub unit_id: String,
}

impl ReductionLogs {
    pub fn empty() -> Self {
        Self {
            evaluations: HashMap::new(),
            events: HashMap::new(),
            exposures: HashMap::new(),
        }
    }

    pub fn evaluation_logs(expression: &Expression) -> Self {
        match expression.is_transient() {
            true => ReductionLogs::empty(),
            false => {
                let id = expression.get_id();
                ReductionLogs {
                    evaluations: HashMap::from([(id.to_string(), 1)]),
                    events: HashMap::new(),
                    exposures: HashMap::new(),
                }
            }
        }
    }

    pub fn exposure_logs(
        split_id: &str,
        unit_id: &str,
        feature_values: &HashMap<String, Value>,
        assignment: &HashMap<String, DbAssignmentEntry>,
    ) -> Result<Self> {
        let exposure = Exposure {
            assignment: assignment.clone(),
            feature_values: feature_values.clone(),
            split_id: split_id.to_string(),
            unit_id: unit_id.to_string(),
        };

        Ok(Self {
            evaluations: HashMap::new(),
            events: HashMap::new(),
            exposures: [(serde_json::to_string(&exposure)?, 1)].into(),
        })
    }

    pub fn event_logs(
        event_type_id: &str,
        unit_id: &str,
        feature_values: &HashMap<String, Value>,
    ) -> Result<Self> {
        let event = Event {
            event_type_id: event_type_id.to_string(),
            feature_values: feature_values.clone(),
            unit_id: unit_id.to_string(),
        };

        Ok(Self {
            evaluations: HashMap::new(),
            events: [(serde_json::to_string(&event)?, 1)].into(),
            exposures: HashMap::new(),
        })
    }

    pub fn merge(a: &Self, b: &Self) -> Self {
        let mut result = a.clone();
        result.evaluations.extend(
            b.evaluations
                .iter()
                .map(|(k, bv)| match a.evaluations.get(k) {
                    Some(av) => (k.clone(), av + bv),
                    None => (k.clone(), *bv),
                }),
        );

        result
            .events
            .extend(b.events.iter().map(|(k, bv)| match a.events.get(k) {
                Some(av) => (k.clone(), av + bv),
                None => (k.clone(), *bv),
            }));

        result
            .exposures
            .extend(b.exposures.iter().map(|(k, bv)| match a.exposures.get(k) {
                Some(av) => (k.clone(), av + bv),
                None => (k.clone(), *bv),
            }));

        result
    }

    pub fn merge_all<'a, T>(logs: T) -> Self
    where
        T: IntoIterator<Item = &'a ReductionLogs>,
    {
        logs.into_iter().fold(ReductionLogs::empty(), |acc, l| {
            ReductionLogs::merge(&acc, l)
        })
    }
}