//! Implements pattern matching

use super::*;
use std::collections::HashMap;

/// Stores a substitution from meta-identifiers to values
#[derive(Clone, Debug)]
pub struct Subst<A: Annotation> {
    pub exp: HashMap<String, Expr<A>>,
    pub pat: HashMap<String, Pat>,
}

impl<A: Annotation> Default for Subst<A> {
    fn default() -> Self {
        Self { exp: Default::default(), pat: Default::default() }
    }
}
impl<A: Annotation> Subst<A> {
    fn add_pat(&mut self, name: String, pat: &Pat) -> bool {
        if self.exp.contains_key(&name) {
            return false;
        }
        if let Some(entry) = self.pat.get(&name) {
            return entry == pat;
        }
        self.pat.insert(name, pat.clone()).is_none()
    }
    fn add_expr(&mut self, name: String, exp: &Expr<A>) -> bool {
        if self.pat.contains_key(&name) {
            return false;
        }
        if let Some(entry) = self.exp.get(&name) {
            return entry == exp;
        }
        self.exp.insert(name, exp.clone()).is_none()
    }
}

pub trait Match<A: Annotation> {
    /// Applies a substitution rule from `pat` to `template` on `self`
    fn apply_rule(&mut self, pat: &Self, template: &Self) -> bool
    where Self: Sized + Clone {
        let Some(sub) = self.construct(pat) else {
            return false;
        };

        *self = template.clone();
        self.apply(&sub);

        true
    }

    /// With self as the pattern, recursively applies the Subst
    fn apply(&mut self, sub: &Subst<A>);

    /// Implements recursive Subst-building for Self
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool;

    /// Constructs a Subst
    fn construct(&self, pat: &Self) -> Option<Subst<A>> {
        let mut sub = Subst::default();
        Match::recurse(&mut sub, pat, self).then_some(sub)
    }

    /// Matches self against the provided pattern
    fn match_with(&self, pat: &Self, sub: &mut Subst<A>) -> bool {
        Match::recurse(sub, pat, self)
    }
}

impl<M: Match<A> + Annotation, A: Annotation> Match<A> for Anno<M, A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        Match::recurse(sub, &pat.0, &expr.0)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        self.0.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Const<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (Self(pat_pat, pat_expr), Self(expr_pat, expr_expr)) = (pat, expr);
        Match::recurse(sub, pat_pat, expr_pat) && Match::recurse(sub, pat_expr, expr_expr)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(pat, expr) = self;
        pat.apply(sub);
        expr.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Fn<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (
            Self(pat_id, pat_arg, _pat_rety, pat_body),
            Self(expr_id, expr_arg, _expr_rety, expr_body),
        ) = (pat, expr);
        pat_id == expr_id
            && Match::recurse(sub, pat_arg, expr_arg)
            && Match::recurse(sub, pat_body, expr_body)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(_, pat, _rety, body) = self;
        pat.apply(sub);
        body.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Let<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (Self(pat_pat, pat_expr), Self(expr_pat, expr_expr)) = (pat, expr);
        Match::recurse(sub, pat_pat, expr_pat) && Match::recurse(sub, pat_expr, expr_expr)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(pat, expr) = self;
        pat.apply(sub);
        expr.apply(sub);
    }
}

impl<A: Annotation> Match<A> for crate::ast::Match<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (Self(pat_expr, pat_arms), Self(expr_expr, expr_arms)) = (pat, expr);
        Match::recurse(sub, pat_expr, expr_expr) && Match::recurse(sub, pat_arms, expr_arms)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(pat, expr) = self;
        pat.apply(sub);
        expr.apply(sub);
    }
}

impl<A: Annotation> Match<A> for crate::ast::Make<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (Make(pat, pat_arms), Make(expr, expr_arms)) = (pat, expr);
        Match::recurse(sub, pat, expr) && Match::recurse(sub, pat_arms, expr_arms)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Make(expr, make_arms) = self;
        expr.apply(sub);
        make_arms.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Mod<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (Self(_pat_ty, pat_expr), Self(_expr_ty, expr_expr)) = (pat, expr);
        // TODO: unify types for ast matching
        // Match::recurse(sub, pat_ty, expr_ty) &&
        Match::recurse(sub, pat_expr, expr_expr)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Mod(_ty, expr) = self;
        expr.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Struct {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        let (Self(pat_pat), Self(expr_pat)) = (pat, expr);
        Match::recurse(sub, pat_pat, expr_pat)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(pat) = self;
        pat.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Expr<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        match (pat, expr) {
            (Expr::MetId(name), _) if name == "_" => true,
            (Expr::MetId(name), _) => sub.add_expr(name.clone(), expr),
            (Expr::Id(pat), Expr::Id(expr)) => pat == expr,
            (Expr::Id(_), _) => false,
            (Expr::Lit(pat), Expr::Lit(expr)) => pat == expr,
            (Expr::Lit(_), _) => false,
            (Expr::Let(pat), Expr::Let(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Let(..), _) => false,
            (Expr::Const(pat), Expr::Const(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Const(..), _) => false,
            (Expr::Struct(pat), Expr::Struct(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Struct(_), _) => false,
            (Expr::Make(pat), Expr::Make(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Make(..), _) => false,
            (Expr::Match(pat), Expr::Match(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Match(..), _) => false,
            (Expr::Mod(pat), Expr::Mod(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Mod(..), _) => false,
            (Expr::Fn(pat), Expr::Fn(expr)) => Match::recurse(sub, pat, expr),
            (Expr::Fn(..), _) => false,
            (Expr::Op(pat_op, pat_exprs), Expr::Op(expr_op, expr_exprs)) => {
                Match::recurse(sub, pat_op, expr_op) && Match::recurse(sub, pat_exprs, expr_exprs)
            }
            (Expr::Op(..), _) => false,
        }
    }

    fn apply(&mut self, sub: &Subst<A>) {
        match self {
            Expr::MetId(id) => {
                if let Some(expr) = sub.exp.get(id) {
                    *self = expr.clone()
                }
            }
            Expr::Id(_) | Expr::Lit(_) => {}
            Expr::Let(expr) => expr.apply(sub),
            Expr::Const(expr) => expr.apply(sub),
            Expr::Struct(expr) => expr.apply(sub),
            Expr::Make(expr) => expr.apply(sub),
            Expr::Match(expr) => expr.apply(sub),
            Expr::Mod(expr) => expr.apply(sub),
            Expr::Fn(expr) => expr.apply(sub),
            Expr::Op(op, exprs) => {
                op.apply(sub);
                exprs.apply(sub);
            }
        };
    }
}

impl<A: Annotation> Match<A> for MakeArm<A> {
    // TODO: order-independent matching for MakeArm specifically.
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        pat.0 == expr.0 && Match::recurse(sub, &pat.1, &expr.1)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(_, expr) = self;
        expr.apply(sub);
    }
}

impl<A: Annotation> Match<A> for MatchArm<A> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        Match::recurse(sub, &pat.0, &expr.0) && Match::recurse(sub, &pat.1, &expr.1)
    }

    fn apply(&mut self, sub: &Subst<A>) {
        let Self(pats, expr) = self;
        pats.apply(sub);
        expr.apply(sub);
    }
}

impl<A: Annotation> Match<A> for Pat {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        match (pat, expr) {
            (Pat::MetId(name), _) if name == "_" => true,
            (Pat::MetId(name), _) => sub.add_pat(name.clone(), expr),
            (Pat::Ignore, Pat::Ignore) => true,
            (Pat::Ignore, _) => false,
            (Pat::Name(pat), Pat::Name(expr)) => pat == expr,
            (Pat::Name(_), _) => false,
            (Pat::Path(_), Pat::Path(_)) => true,
            (Pat::Path(_), _) => false,
            (Pat::Struct(_, pat), Pat::Struct(_, expr)) => Match::recurse(sub, pat, expr),
            (Pat::Struct(..), _) => false,
            (Pat::TupStruct(_, pat), Pat::TupStruct(_, expr)) => Match::recurse(sub, pat, expr),
            (Pat::TupStruct(..), _) => false,
            (Pat::Lit(pat), Pat::Lit(expr)) => pat == expr,
            (Pat::Lit(_), _) => false,
            (Pat::Typed(pat, _), Pat::Typed(expr, _)) => Match::recurse(sub, pat, expr),
            (Pat::Typed(..), _) => false,
            (Pat::Op(_, pat), Pat::Op(_, expr)) => Match::recurse(sub, pat, expr),
            (Pat::Op(..), _) => false,
        }
    }

    fn apply(&mut self, sub: &Subst<A>) {
        match self {
            Pat::Ignore | Pat::Name(_) | Pat::Path(_) | Pat::Lit(_) => {}
            Pat::MetId(id) => {
                if let Some(expr) = sub.pat.get(id) {
                    *self = expr.clone()
                }
            }
            Pat::Struct(_, expr) => expr.apply(sub),
            Pat::TupStruct(_, expr) => expr.apply(sub),
            Pat::Typed(pat, ty) => {
                pat.apply(sub);
                ty.apply(sub);
            }
            Pat::Op(_, pats) => pats.apply(sub),
        }
    }
}

impl<A: Annotation> Match<A> for Ty {
    fn apply(&mut self, sub: &Subst<A>) {
        todo!("Apply subst {sub:?} for {self}.")
    }

    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        todo!("Construct subst {sub:?} from {pat} and {expr}.")
    }
}

impl<A: Annotation> Match<A> for Op {
    fn recurse(_: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        pat == expr
    }

    fn apply(&mut self, _sub: &Subst<A>) {}
}

impl<A: Annotation, T: Match<A>> Match<A> for [T] {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        if pat.len() != expr.len() {
            return false;
        }
        for (pat, expr) in pat.iter().zip(expr.iter()) {
            if !Match::recurse(sub, pat, expr) {
                return false;
            }
        }
        true
    }

    fn apply(&mut self, sub: &Subst<A>) {
        for item in self {
            item.apply(sub);
        }
    }
}

impl<A: Annotation, T: Match<A>> Match<A> for Box<T> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        Match::recurse(sub, pat.as_ref(), expr.as_ref())
    }

    fn apply(&mut self, sub: &Subst<A>) {
        self.as_mut().apply(sub);
    }
}

impl<A: Annotation, T: Match<A>> Match<A> for Vec<T> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        Match::recurse(sub, pat.as_slice(), expr.as_slice())
    }

    fn apply(&mut self, sub: &Subst<A>) {
        self.as_mut_slice().apply(sub);
    }
}

impl<A: Annotation, T: Match<A>> Match<A> for Option<T> {
    fn recurse(sub: &mut Subst<A>, pat: &Self, expr: &Self) -> bool {
        match (pat, expr) {
            (Some(pat), Some(expr)) => Match::recurse(sub, pat, expr),
            (None, None) => true,
            _ => false,
        }
    }

    fn apply(&mut self, sub: &Subst<A>) {
        self.as_mut_slice().apply(sub);
    }
}