/*
 * Minimal term re-writing system
 */

use crate::cexpr::{CExpr, CNumber};
use crate::sexpr::SExpr;
use crate::Result;
use std::collections::HashMap;

// for now these are hard coded
#[derive(Clone, PartialEq, Debug)]
pub enum RulePredicate {
    IsNumber,
    IsSymbol,
    IsInteger,
    IsEven,
    IsOdd,
    IsPositive,
    IsNegative,
}

#[derive(Clone, PartialEq, Debug)]
pub enum CExprType {
    Sum,
    Product,
    Factorial,
    Exponential,
    UnaryFunction(String),
}

#[derive(Clone, PartialEq, Debug)]
pub enum RuleExpr {
    Var(char),
    VarList(char),
    Atom(CExpr),
    List(CExprType, Vec<RuleExpr>),
}

#[derive(Clone, PartialEq, Debug)]
pub struct Rule {
    variables: HashMap<char, Vec<RulePredicate>>,
    pattern: RuleExpr,
    substitute: RuleExpr,
}

impl RulePredicate {
    pub fn from_sexpr(sexpr: &SExpr) -> Result<RulePredicate> {
        use RulePredicate::*;
        match sexpr {
            SExpr::SIdentifier(v) => match v.as_str() {
                "number?" => Ok(IsNumber),
                "symbol?" => Ok(IsSymbol),
                "integer?" => Ok(IsInteger),
                "even?" => Ok(IsEven),
                "odd?" => Ok(IsOdd),
                "positive?" => Ok(IsPositive),
                "negative?" => Ok(IsNegative),
                _ => Err(format!("not a known predicate: {:?}", sexpr)),
            },
            _ => Err(format!("not a known predicate: {:?}", sexpr)),
        }
    }

    pub fn check(self, cexpr: &CExpr) -> bool {
        use CExpr::*;
        use CNumber::*;
        use RulePredicate::*;
        match (self, cexpr) {
            (IsNumber, Number(_)) => true,
            (IsSymbol, Symbol(_)) => true,
            (IsInteger, Number(Integer(_))) => true,
            (IsEven, Number(Integer(n))) => n % 2 == 0,
            (IsOdd, Number(Integer(n))) => n % 2 == 1,
            (IsPositive, Number(Integer(n))) => n >= &0,
            (IsNegative, Number(Integer(n))) => n < &0,
            _ => false,
        }
    }
}

fn extract_vars(sexpr: &SExpr) -> Result<Option<HashMap<char, Vec<RulePredicate>>>> {
    let slist = match sexpr {
        SExpr::SList(l) => l,
        _ => return Ok(None),
    };
    if matches!(slist.as_slice(), [SExpr::SIdentifier(head), SExpr::SIdentifier(var_name), ..] if head == "?" && var_name.len() == 1)
    {
        if let [SExpr::SIdentifier(_), SExpr::SIdentifier(var_name), pred_list @ ..] =
            slist.as_slice()
        {
            let predicates: Vec<RulePredicate> = pred_list
                .iter()
                .map(|v| RulePredicate::from_sexpr(v))
                .collect::<Result<Vec<RulePredicate>>>()?;
            // TODO: why did HashMap::from([(c, predicates)]) not work here? rustc version?
            let mut hm = HashMap::new();
            hm.insert(var_name.chars().nth(0).unwrap(), predicates);
            return Ok(Some(hm));
        }
    }
    let sub_vars = slist.iter().filter_map(|v| extract_vars(v).unwrap()).fold(
        HashMap::new(),
        |mut acc, hm| {
            acc.extend(hm);
            acc
        },
    );
    if sub_vars.is_empty() {
        return Ok(None);
    } else {
        return Ok(Some(sub_vars));
    }
}

impl RuleExpr {
    pub fn from_sexpr(sexpr: &SExpr) -> Result<RuleExpr> {
        // not all cases are handled; some atoms are covered trivialy
        match sexpr {
            SExpr::SNull | SExpr::SBoolean(_) | SExpr::SString(_) => {
                Err("unhandled s-expr atoms; expected rule pattern".to_string())
            }
            SExpr::SInteger(_) | SExpr::SFloat(_) | SExpr::SIdentifier(_) => {
                Ok(RuleExpr::Atom(CExpr::from_sexpr(sexpr)?))
            }
            SExpr::SList(list) => {
                if let [SExpr::SIdentifier(ident), rest @ ..] = list.as_slice() {
                    match (ident.as_str(), rest.len()) {
                        ("?", 1..=5000) => {
                            if let SExpr::SIdentifier(var_name) = &rest[0] {
                                Ok(RuleExpr::Var(var_name.chars().nth(0).unwrap()))
                            } else {
                                Err("expected single-character pattern name".to_string())
                            }
                        }
                        ("?*", 1..=5000) => {
                            if let SExpr::SIdentifier(var_name) = &rest[0] {
                                Ok(RuleExpr::VarList(var_name.chars().nth(0).unwrap()))
                            } else {
                                Err("expected single-character pattern name".to_string())
                            }
                        }
                        ("factorial", 1) => Ok(RuleExpr::List(
                            CExprType::Factorial,
                            rest.iter()
                                .map(|v| RuleExpr::from_sexpr(v))
                                .collect::<Result<Vec<RuleExpr>>>()?,
                        )),
                        ("cos" | "sin" | "tan", 1) => Ok(RuleExpr::List(
                            CExprType::UnaryFunction(ident.to_string()),
                            rest.iter()
                                .map(|v| RuleExpr::from_sexpr(v))
                                .collect::<Result<Vec<RuleExpr>>>()?,
                        )),
                        ("^", 2) => Ok(RuleExpr::List(
                            CExprType::Exponential,
                            rest.iter()
                                .map(|v| RuleExpr::from_sexpr(v))
                                .collect::<Result<Vec<RuleExpr>>>()?,
                        )),
                        ("+", 2..=5000) => Ok(RuleExpr::List(
                            CExprType::Sum,
                            rest.iter()
                                .map(|v| RuleExpr::from_sexpr(v))
                                .collect::<Result<Vec<RuleExpr>>>()?,
                        )),
                        ("*", 2..=5000) => Ok(RuleExpr::List(
                            CExprType::Product,
                            rest.iter()
                                .map(|v| RuleExpr::from_sexpr(v))
                                .collect::<Result<Vec<RuleExpr>>>()?,
                        )),
                        _ => Err("unhandled rule expression type".to_string()),
                    }
                } else {
                    Err("unhandled rule expression type".to_string())
                }
            }
        }
    }
}

impl Rule {
    pub fn from_sexpr(sexpr: &SExpr) -> Result<Rule> {
        let slist = match sexpr {
            SExpr::SList(l) => l,
            _ => return Err("expected a rule".to_string()),
        };
        let (pattern, substitute) =
            if let [SExpr::SIdentifier(name), pattern, substitute] = slist.as_slice() {
                if name != "rule" {
                    return Err("expected a rule".to_string());
                }
                (pattern, substitute)
            } else {
                return Err("expected a rule".to_string());
            };

        // extract vars and predicates from pattern
        let vars = match extract_vars(&pattern)? {
            Some(hm) => hm,
            None => HashMap::new(),
        };

        // parse pattern and substitute as RuleExpr
        Ok(Rule {
            variables: vars,
            pattern: RuleExpr::from_sexpr(&pattern)?,
            substitute: RuleExpr::from_sexpr(&substitute)?,
        })
    }

    pub fn from_str(raw: &str) -> Result<Rule> {
        let ast = SExpr::from_str(raw)?;
        Rule::from_sexpr(&ast)
    }
}