use std::collections::HashMap;

use crate::sexpr::{Atom, Expr, Variable, VariableType};

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum BindValue {
    Integer(i64),
    Expr(Expr),
    Ellipsis(Vec<Expr>)
}

#[derive(Default)]
pub struct Bindings<'b>(HashMap<&'b str, BindValue>);

impl<'b> Bindings<'b> {
    pub fn check_or_insert(&mut self, index: &'b String, bind_value: BindValue) -> bool {
        if let Some(existing) = self.0.get(index.as_str()) {
            existing == &bind_value
        } else {
            self.0.insert(index, bind_value);
            true
        }
    }
}

impl std::fmt::Debug for Bindings<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self.0)
    }
}

pub fn args_match<'b>(p_args: &'b [Expr], args: &'b [Expr], bindings: &mut Bindings<'b>) -> bool {
    const ELLIPSIS: Expr = Expr::Atom(Atom::Ellipsis);

    if p_args.contains(&ELLIPSIS) {
        let mut sorted_p: Vec<&Expr> = p_args.iter()
            .filter(|p| **p != ELLIPSIS)
            .collect();
        
        let mut sorted_a: Vec<&Expr> = args.iter().collect();

        if sorted_p.len() != sorted_a.len() {
            return false;
        }

        sorted_p.sort(); 
        sorted_a.sort();

        std::iter::zip(sorted_p, sorted_a).all(|(p, a)| {
            matches(p, a, bindings)
        })
    } else {
        if p_args.len() != args.len() {
            return false;
        }
        std::iter::zip(p_args.iter(), args.iter()).all(|(a, b)| {
            matches(a, b, bindings)
        })
    }
}

pub fn matches<'b>(p: &'b Expr, expr: &'b Expr, bindings: &mut Bindings<'b>) -> bool {
    match expr {
        Expr::Atom(Atom::Int(i)) => {
            match p {
                Expr::Atom(Atom::Int(pi)) => i == pi,
                Expr::Atom(Atom::Variable(Variable { r#type: e @ (VariableType::Integer | VariableType::Expr), index })) => {
                    bindings.check_or_insert(index, match e {
                        VariableType::Integer => BindValue::Integer(*i),
                        VariableType::Expr => BindValue::Expr(Expr::Atom(Atom::Int(*i)))
                    })
                }
                _ => false
            }
        }
        Expr::Atom(Atom::Builtin(s)) => {
            match p {
                Expr::Atom(Atom::Builtin(ps)) => s == ps,
                Expr::Atom(Atom::Variable(Variable { r#type: VariableType::Expr, index })) => {
                    bindings.check_or_insert(index, BindValue::Expr(Expr::Atom(Atom::Builtin(s.to_string()))))
                }
                _ => false
            }
        }
        Expr::Application(exprs) => {
            if let Expr::Atom(Atom::Builtin(f)) = &exprs[0] &&
               let Expr::Application(p_exprs) = p &&
               let Expr::Atom(Atom::Builtin(pf)) = &p_exprs[0] {
                f == pf && args_match(&p_exprs[1..], &exprs[1..], bindings)
            } else if exprs.len() == 1 {
                matches(p, &exprs[0], bindings)
            } else if let Expr::Atom(Atom::Variable(Variable { r#type: VariableType::Expr, index })) = p {
                bindings.check_or_insert(index, BindValue::Expr(expr.clone()))
            } else {
                false
            }
        },
        Expr::Atom(Atom::Variable(_)) | Expr::Atom(Atom::Ellipsis) => {
            // These should only appear in patterns, not in expressions being matched
            false
        }
    }
}