use crate::base::*;
use crate::parser::*;
use crate::value::{Fn, Value};
use imbl::HashMap;
use imbl::Vector;
use std::cell::RefCell;
use std::rc::Rc;

pub fn match_eq<T, U>(x: T, y: T, z: U) -> Option<U>
where
    T: PartialEq,
{
    if x == y {
        Some(z)
    } else {
        None
    }
}

pub struct LErr {
    pub msg: String,
    pub trace: Vec<String>,
}

impl LErr {
    pub fn new(msg: String) -> LErr {
        LErr {
            msg: msg.to_string(),
            trace: vec![],
        }
    }

    pub fn add_trace(mut self, fn_name: String) -> LErr {
        self.trace.push(fn_name);
        self
    }
}

type LResult<'src> = Result<Value<'src>, LErr>;

pub struct Context<'src> {
    pub locals: Vec<(String, Value<'src>)>,
    pub prelude: Vec<(String, Value<'src>)>,
    pub prelude_ast: &'src Ast,
    pub ast: &'src Ast,
}

impl<'src> Context<'src> {
    pub fn resolve(&self, word: &String) -> LResult<'src> {
        let resolved_local = self.locals.iter().rev().find(|(name, _)| word == name);

        match resolved_local {
            Some((_, value)) => Ok(value.clone()),
            None => {
                let resolved_prelude = self.prelude.iter().rev().find(|(name, _)| word == name);
                match resolved_prelude {
                    Some((_, value)) => Ok(value.clone()),
                    None => Err(LErr::new(format!("unbound name {word}"))),
                }
            }
        }
    }

    pub fn bind(&mut self, word: String, value: &Value<'src>) {
        self.locals.push((word, value.clone()));
    }

    pub fn pop_to(&mut self, n: usize) {
        while self.locals.len() > n {
            self.locals.pop();
        }
    }

    pub fn match_eq<T>(&self, x: T, y: T) -> Option<&Context<'src>>
    where
        T: PartialEq,
    {
        if x == y {
            Some(&self)
        } else {
            None
        }
    }

    pub fn match_pattern(&mut self, patt: &Pattern, val: &Value<'src>) -> Option<&Context<'src>> {
        match (patt, val) {
            (Pattern::Nil, Value::Nil) => Some(self),
            (Pattern::Placeholder, _) => Some(self),
            (Pattern::Number(x), Value::Number(y)) => self.match_eq(x, y),
            (Pattern::Boolean(x), Value::Boolean(y)) => self.match_eq(x, y),
            (Pattern::Keyword(x), Value::Keyword(y)) => self.match_eq(x, y),
            (Pattern::String(x), Value::InternedString(y)) => self.match_eq(x, y),
            (Pattern::String(x), Value::AllocatedString(y)) => self.match_eq(&x.to_string(), y),
            (Pattern::Interpolated(_, StringMatcher(matcher)), Value::InternedString(y)) => {
                match matcher(y.to_string()) {
                    Some(matches) => {
                        let mut matches = matches
                            .iter()
                            .map(|(word, string)| {
                                (
                                    word.clone(),
                                    Value::AllocatedString(Rc::new(string.clone())),
                                )
                            })
                            .collect::<Vec<_>>();
                        self.locals.append(&mut matches);
                        Some(self)
                    }
                    None => None,
                }
            }
            (Pattern::Word(w), val) => {
                self.bind(w.to_string(), &val);
                Some(self)
            }
            (Pattern::As(word, type_str), value) => {
                let ludus_type = r#type(value);
                let type_kw = Value::Keyword(type_str);
                if type_kw == ludus_type {
                    self.bind(word.to_string(), value);
                    Some(self)
                } else {
                    None
                }
            }
            (Pattern::Tuple(x), Value::Tuple(y)) => {
                let has_splat = x
                    .iter()
                    .any(|patt| matches!(patt, (Pattern::Splattern(_), _)));
                if x.len() > y.len() || (!has_splat && x.len() != y.len()) {
                    return None;
                };
                let to = self.locals.len();
                for i in 0..x.len() {
                    if let Pattern::Splattern(patt) = &x[i].0 {
                        let mut list = Vector::new();
                        for i in i..y.len() {
                            list.push_back(y[i].clone())
                        }
                        let list = Value::List(list);
                        self.match_pattern(&patt.0, &list);
                    } else if self.match_pattern(&x[i].0, &y[i]).is_none() {
                        self.pop_to(to);
                        return None;
                    }
                }
                Some(self)
            }
            (Pattern::List(x), Value::List(y)) => {
                let has_splat = x
                    .iter()
                    .any(|patt| matches!(patt, (Pattern::Splattern(_), _)));
                if x.len() > y.len() || (!has_splat && x.len() != y.len()) {
                    return None;
                };
                let to = self.locals.len();
                for (i, (patt, _)) in x.iter().enumerate() {
                    if let Pattern::Splattern(patt) = &patt {
                        let list = Value::List(y.skip(i));
                        self.match_pattern(&patt.0, &list);
                    } else if self.match_pattern(patt, y.get(i).unwrap()).is_none() {
                        self.pop_to(to);
                        return None;
                    }
                }
                Some(self)
            }
            // TODO: optimize this on several levels
            // - [ ] opportunistic mutation
            // - [ ] get rid of all the pointer indirection in word splats
            (Pattern::Dict(x), Value::Dict(y)) => {
                let has_splat = x
                    .iter()
                    .any(|patt| matches!(patt, (Pattern::Splattern(_), _)));
                if x.len() > y.len() || (!has_splat && x.len() != y.len()) {
                    return None;
                };
                let to = self.locals.len();
                let mut matched = vec![];
                for (pattern, _) in x {
                    match pattern {
                        Pattern::Pair(key, patt) => {
                            if let Some(val) = y.get(key) {
                                if self.match_pattern(&patt.0, val).is_none() {
                                    self.pop_to(to);
                                    return None;
                                } else {
                                    matched.push(key);
                                }
                            } else {
                                return None;
                            };
                        }
                        Pattern::Splattern(pattern) => match pattern.0 {
                            Pattern::Word(w) => {
                                // TODO: find a way to take ownership
                                // this will ALWAYS make structural changes, because of this clone
                                // we want opportunistic mutation if possible
                                let mut unmatched = y.clone();
                                for key in matched.iter() {
                                    unmatched.remove(*key);
                                }
                                self.bind(w.to_string(), &Value::Dict(unmatched));
                            }
                            Pattern::Placeholder => (),
                            _ => unreachable!(),
                        },
                        _ => unreachable!(),
                    }
                }
                Some(self)
            }
            _ => None,
        }
    }

    pub fn match_clauses(
        &mut self,
        value: &Value<'src>,
        clauses: &'src [MatchClause],
    ) -> LResult<'src> {
        {
            let parent = self.ast;
            let to = self.locals.len();
            for MatchClause { patt, body, guard } in clauses.iter() {
                if let Some(_) = self.match_pattern(&patt.0, value) {
                    let pass_guard = match guard {
                        None => true,
                        Some((ast, _)) => {
                            self.ast = ast;
                            let guard_res = self.eval();
                            match &guard_res {
                                Err(_) => return guard_res,
                                Ok(val) => val.bool(),
                            }
                        }
                    };
                    if !pass_guard {
                        self.pop_to(to);
                        continue;
                    }
                    self.ast = &body.0;
                    let res = self.eval();
                    self.pop_to(to);
                    self.ast = parent;
                    return res;
                }
            }
            Err(LErr::new(format!("no match")))
        }
    }

    pub fn apply(&mut self, callee: Value<'src>, caller: Value<'src>) -> LResult<'src> {
        match (callee, caller) {
            (Value::Keyword(kw), Value::Dict(dict)) => {
                if let Some(val) = dict.get(kw) {
                    Ok(val.clone())
                } else {
                    Ok(Value::Nil)
                }
            }
            (Value::Dict(dict), Value::Keyword(kw)) => {
                if let Some(val) = dict.get(kw) {
                    Ok(val.clone())
                } else {
                    Ok(Value::Nil)
                }
            }
            (Value::Fn(f), Value::Tuple(args)) => {
                let args = Value::Tuple(args);
                self.match_clauses(&args, f.body)
            }
            (Value::Fn(_f), Value::Args(_args)) => todo!(),
            (_, Value::Keyword(_)) => Ok(Value::Nil),
            (_, Value::Args(_)) => Err(LErr::new(format!("you may only call a function"))),
            (Value::Base(f), Value::Tuple(args)) => match f {
                Base::Nullary(f) => {
                    let num_args = args.len();
                    if num_args != 0 {
                        Err(LErr::new(format!(
                            "wrong arity: expected 0 arguments, got {num_args}"
                        )))
                    } else {
                        Ok(f())
                    }
                }
                Base::Unary(f) => {
                    let num_args = args.len();
                    if num_args != 1 {
                        Err(LErr::new(format!(
                            "wrong arity: expected 1 argument, got {num_args}"
                        )))
                    } else {
                        Ok(f(&args[0]))
                    }
                }
                Base::Binary(r#fn) => {
                    let num_args = args.len();
                    if num_args != 2 {
                        Err(LErr::new(format!(
                            "wrong arity: expected 2 arguments, got {num_args}"
                        )))
                    } else {
                        Ok(r#fn(&args[0], &args[1]))
                    }
                }
                Base::Ternary(f) => {
                    let num_args = args.len();
                    if num_args != 3 {
                        Err(LErr::new(format!(
                            "wrong arity: expected 3 arguments, got {num_args}"
                        )))
                    } else {
                        Ok(f(&args[0], &args[1], &args[2]))
                    }
                }
            },
            _ => unreachable!(),
        }
    }

    pub fn eval(&mut self) -> LResult<'src> {
        let root = self.ast;
        let result = match self.ast {
            Ast::Nil => Ok(Value::Nil),
            Ast::Boolean(b) => Ok(Value::Boolean(*b)),
            Ast::Number(n) => Ok(Value::Number(*n)),
            Ast::Keyword(k) => Ok(Value::Keyword(k)),
            Ast::String(s) => Ok(Value::InternedString(s)),
            Ast::Interpolated(parts) => {
                let mut interpolated = String::new();
                for part in parts {
                    match &part.0 {
                        StringPart::Data(s) => interpolated.push_str(s.as_str()),
                        StringPart::Word(w) => {
                            let val = self.resolve(w)?;
                            interpolated.push_str(val.interpolate().as_str())
                        }
                        StringPart::Inline(_) => unreachable!(),
                    }
                }
                Ok(Value::AllocatedString(Rc::new(interpolated)))
            }
            Ast::Block(exprs) => {
                let parent = self.ast;
                let to = self.locals.len();
                let mut result = Value::Nil;
                for (expr, _) in exprs {
                    self.ast = &expr;
                    result = self.eval()?;
                }
                self.pop_to(to);
                self.ast = parent;
                Ok(result)
            }
            Ast::If(cond, if_true, if_false) => {
                let parent = self.ast;
                self.ast = &cond.0;
                let truthy = self.eval()?.bool();
                self.ast = if truthy { &if_true.0 } else { &if_false.0 };
                let result = self.eval();
                self.ast = parent;
                result
            }
            Ast::List(members) => {
                let parent = self.ast;
                let mut vect = Vector::new();
                for member in members {
                    self.ast = &member.0;
                    if let Ast::Splat(_) = self.ast {
                        let to_splat = self.eval()?;
                        match to_splat {
                            Value::List(list) => vect.append(list),
                            _ => {
                                return Err(LErr::new(format!(
                                    "only lists may be splatted into lists"
                                )))
                            }
                        }
                    } else {
                        vect.push_back(self.eval()?)
                    }
                }
                self.ast = parent;
                Ok(Value::List(vect))
            }
            Ast::Tuple(members) => {
                let parent = self.ast;
                let mut vect = Vec::new();
                for member in members {
                    self.ast = &member.0;
                    vect.push(self.eval()?);
                }
                self.ast = parent;
                Ok(Value::Tuple(Rc::new(vect)))
            }
            Ast::Word(w) | Ast::Splat(w) => {
                let val = self.resolve(&w.to_string())?;
                Ok(val)
            }
            Ast::Let(patt, expr) => {
                let parent = self.ast;
                self.ast = &expr.0;
                let val = self.eval()?;
                let result = match self.match_pattern(&patt.0, &val) {
                    Some(_) => Ok(val),
                    None => Err(LErr::new(format!("no match"))),
                };
                self.ast = parent;
                result
            }
            Ast::Placeholder => Ok(Value::Placeholder),
            Ast::Error => unreachable!(),
            Ast::Arguments(a) => {
                let parent = self.ast;
                let mut args = vec![];
                for (arg, _) in a.iter() {
                    self.ast = arg;
                    let arg = self.eval()?;
                    args.push(arg);
                }
                let result = if args.iter().any(|arg| matches!(arg, Value::Placeholder)) {
                    Ok(Value::Args(Rc::new(args)))
                } else {
                    Ok(Value::Tuple(Rc::new(args)))
                };
                self.ast = parent;
                result
            }
            Ast::Dict(terms) => {
                let parent = self.ast;
                let mut dict = HashMap::new();
                for term in terms {
                    let (term, _) = term;
                    match term {
                        Ast::Pair(key, value) => {
                            self.ast = &value.0;
                            let value = self.eval()?;
                            dict.insert(*key, value);
                        }
                        Ast::Splat(_) => {
                            self.ast = term;
                            let resolved = self.eval()?;
                            let Value::Dict(to_splat) = resolved else {
                                return Err(LErr::new(format!("cannot splat non-dict into dict")));
                            };
                            dict = to_splat.union(dict);
                        }
                        _ => unreachable!(),
                    }
                }
                self.ast = parent;
                Ok(Value::Dict(dict))
            }
            Ast::Box(name, expr) => {
                let parent = self.ast;
                self.ast = &expr.0;
                let val = self.eval()?;
                let boxed = Value::Box(name, Rc::new(RefCell::new(val)));
                self.bind(name.to_string(), &boxed);
                self.ast = parent;
                Ok(boxed)
            }
            Ast::Synthetic(root, first, rest) => {
                let parent = self.ast;
                self.ast = &root.0;
                let root = self.eval()?;
                self.ast = &first.0;
                let first = self.eval()?;
                let mut curr = self.apply(root, first)?;
                for term in rest.iter() {
                    self.ast = &term.0;
                    let next = self.eval()?;
                    curr = self.apply(curr, next)?;
                }
                self.ast = parent;
                Ok(curr)
            }
            Ast::When(clauses) => {
                let parent = self.ast;
                for clause in clauses.iter() {
                    let WhenClause { cond, body } = &clause.0;
                    self.ast = &cond.0;
                    if self.eval()?.bool() {
                        self.ast = &body.0;
                        let result = self.eval();
                        self.ast = parent;
                        return result;
                    };
                }
                Err(LErr::new(format!("no match")))
            }
            Ast::Match(value, clauses) => {
                let parent = self.ast;
                self.ast = &value.0;
                let value = self.eval()?;
                let result = self.match_clauses(&value, clauses);
                self.ast = parent;
                result
            }
            Ast::Fn(name, clauses, doc) => {
                let doc = doc.map(|s| s.to_string());
                let the_fn = Value::Fn::<'src>(Rc::new(Fn::<'src> {
                    name: name.to_string(),
                    body: clauses,
                    doc,
                }));
                self.bind(name.to_string(), &the_fn);
                Ok(the_fn)
            }
            Ast::FnDeclaration(_name) => Ok(Value::Nil),
            Ast::Panic(msg) => {
                self.ast = &msg.0;
                let msg = self.eval()?;
                Err(LErr::new(format!("{msg}")))
            }
            Ast::Repeat(times, body) => {
                let parent = self.ast;
                self.ast = &times.0;
                let times_num = match self.eval() {
                    Ok(Value::Number(n)) => n as usize,
                    _ => return Err(LErr::new(format!("repeat may only take numbers"))),
                };
                self.ast = &body.0;
                for _ in 0..times_num {
                    self.eval()?;
                }
                self.ast = parent;
                Ok(Value::Nil)
            }
            Ast::Do(terms) => {
                let parent = self.ast;
                self.ast = &terms[0].0;
                let mut result = self.eval()?;
                for (term, _) in terms.iter().skip(1) {
                    self.ast = term;
                    let next = self.eval()?;
                    let arg = Value::Tuple(Rc::new(vec![result]));
                    result = self.apply(next, arg)?;
                }
                self.ast = parent;
                Ok(result)
            }
            Ast::Pair(..) => {
                unreachable!()
            }
            Ast::Loop(init, clauses) => {
                let parent = self.ast;
                self.ast = &init.0;
                let mut args = self.eval()?;
                loop {
                    let result = self.match_clauses(&args, clauses)?;
                    if let Value::Recur(recur_args) = result {
                        args = Value::Tuple(Rc::new(recur_args));
                    } else {
                        self.ast = parent;
                        return Ok(result);
                    }
                }
            }
            Ast::Recur(args) => {
                let parent = self.ast;
                let mut vect = Vec::new();
                for arg in args {
                    self.ast = &arg.0;
                    vect.push(self.eval()?);
                }
                self.ast = parent;
                Ok(Value::Recur(vect))
            }
        };
        self.ast = root;
        result
    }
}