#![feature(let_else)]
#![feature(array_windows)]
#![feature(box_patterns)]

mod grammar;

#[cfg(test)]
mod tests;


// use debug2::dbg;
use rustyline::validate::{
    MatchingBracketValidator, ValidationContext, ValidationResult, Validator,
};
use rustyline::Editor;
use rustyline_derive::{Completer, Helper, Highlighter, Hinter};
use std::collections::HashMap;
use std::io;
use std::rc::Rc;

#[derive(Debug, debug2::Debug, PartialEq)]
enum Tree {
    Leaf(Literal),
    Define(String, Box<Tree>),
    If(Box<[Tree; 3]>),
    Lambda(Lambda),
    Branch(Vec<Tree>),
}

#[derive(Debug, debug2::Debug, PartialEq, Clone)]
struct Lambda(Rc<[String]>, Rc<[Tree]>);

#[derive(Debug, debug2::Debug, PartialEq)]
enum Literal {
    Sym(String),
    Num(f64),
    Bool(bool),
}

#[derive(Debug)]
struct RTError(String);

fn err<T>(s: String) -> Result<T, RTError> {
    Err(RTError(s))
}

type Func = fn(&[Value]) -> Result<Value, RTError>;

#[derive(Clone)]
enum Value {
    Num(f64),
    Func(Func),
    Lambda(Lambda),
    Bool(bool),
    NotAValue,
}

impl PartialEq for Value {
    fn eq(&self, other: &Self) -> bool {
        use Value::*;
        match (self, other) {
            (Num(l), Num(r)) => l == r,
            (Func(l), Func(r)) => *l as usize == *r as usize,
            (Bool(l), Bool(r)) => l == r,
            (Lambda(l), Lambda(r)) => Rc::ptr_eq(&l.0, &r.0) && Rc::ptr_eq(&l.1, &r.1),

            (Num(_), _) => false,
            (Func(_), _) => false,
            (Bool(_), _) => false,
            (Lambda(_), _) => false,

            (NotAValue, _) => panic!("Trap value"),
        }
    }
}

impl std::fmt::Debug for Value {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Num(n) => n.fmt(f),
            Self::Func(_) => f.write_str("#<procedure>"),
            Self::Lambda(_) => f.write_str("#<procedure>"),
            Self::Bool(b) => f.write_str(if *b { "#t" } else { "#f" }),
            Self::NotAValue => Ok(()),
        }
    }
}

impl Value {
    fn as_num(&self) -> Result<f64, RTError> {
        if let Self::Num(n) = self {
            Ok(*n)
        } else {
            Err(RTError("Expected a number".to_owned()))
        }
    }

    fn as_func(&self) -> Result<Callable, RTError> {
        match self {
            Self::Func(f) => Ok(Callable::Func(*f)),
            Self::Lambda(l) => Ok(Callable::Lambda(l)),
            _ => Err(RTError(format!("Expected a function, got {:?}", self))),
        }
    }

    fn as_bool(&self) -> Result<bool, RTError> {
        if let Self::Bool(b) = self {
            Ok(*b)
        } else {
            Err(RTError("Expected a bool".to_owned()))
        }
    }
}

enum Callable<'a> {
    Func(Func),
    Lambda(&'a Lambda),
}

struct Env<'a> {
    vars: HashMap<String, Value>,
    enclosing: Option<&'a Env<'a>>,
}

impl<'a> Env<'a> {
    fn child(&'a self) -> Env<'a> {
        Env {
            vars: HashMap::new(),
            enclosing: Some(self),
        }
    }

    fn lookup(&self, s: &str) -> Option<Value> {
        if let Some(v) = self.vars.get(s) {
            Some(v.clone())
        } else if let Some(parent) = self.enclosing {
            parent.lookup(s)
        } else {
            None
        }
    }

    fn define(&mut self, name: String, val: Value) {
        self.vars.insert(name, val);
    }
}

fn main() {
    let args = std::env::args().collect::<Vec<_>>();
    match &args[..] {
        [_] => repl(),
        [_, file] => run_file(file).unwrap(),
        _ => panic!("To many args `{:?}`", args),
    }
}

fn run_file(file: &str) -> io::Result<()> {
    let src = std::fs::read_to_string(file)?;
    let tree = grammar::TreesParser::new().parse(&src).unwrap();
    let mut env = default_env();
    for i in tree {
        eval(&i, &mut env).unwrap();
    }
    Ok(())
}

fn repl() {
    let mut rl = Editor::new();
    rl.set_helper(Some(InputValidator {
        brackets: MatchingBracketValidator::new(),
    }));

    let mut env = default_env();

    while let Ok(line) = rl.readline("> ") {
        rl.add_history_entry(&line);

        let tree = grammar::TreeParser::new().parse(&line).unwrap();
        // dbg!(&tree);
        println!("< {:?}", eval(&tree, &mut env))
    }
}

fn eval(t: &Tree, env: &mut Env) -> Result<Value, RTError> {
    Ok(match t {
        Tree::Leaf(l) => match l {
            Literal::Sym(s) => match env.lookup(s) {
                Some(v) => v.clone(),
                None => return err(format!("Undefined variable `{}`", s)),
            },
            Literal::Num(v) => Value::Num(*v),
            Literal::Bool(b) => Value::Bool(*b),
        },
        Tree::Lambda(l) => Value::Lambda(l.clone()),
        Tree::Branch(args) => {
            let Some(fun) = args.get(0) else { return err("No argument given".to_owned()) };
            let fun = eval(fun, env)?;
            let fun = fun.as_func()?;
            let args = args
                .iter()
                .skip(1)
                .map(|a| eval(a, env))
                .collect::<Result<Vec<_>, _>>()?;
            // return fun(&args);
            match fun {
                Callable::Func(f) => f(&args)?,
                Callable::Lambda(l) => {
                    if l.0.len() == args.len() {
                        let mut env = env.child();
                        for (x, y) in l.0.iter().zip(args) {
                            env.define(x.clone(), y)
                        }

                        let [main @ .., tail] = &l.1[..] else { unreachable!("Body has 1+ element by parser") };
                        for i in main {
                            eval(i, &mut env)?;
                        }

                        eval(tail, &mut env)?
                    } else {
                        return err(format!("Need {} args, got {}", l.0.len(), args.len()));
                    }
                }
            }
        }
        Tree::Define(name, to) => {
            let val = eval(to, env)?;
            env.define(name.to_owned(), val);
            Value::NotAValue
        }
        Tree::If(box [cond, tcase, fcase]) => {
            let b = eval(cond, env)?.as_bool()?;
            let body = if b { tcase } else { fcase };
            eval(body, env)?
        }
    })
}

fn default_env() -> Env<'static> {
    let mut vars = HashMap::new();

    for (name, fun) in [
        ("*", prims::mul as Func),
        ("+", prims::add as Func),
        ("-", prims::sub as Func),
        ("/", prims::div as Func),
        // TODO: Spec compiance
        ("=", prims::equals as Func),
        ("eq?", prims::equals as Func),
        ("eqv?", prims::equals as Func),
        ("equal?", prims::equals as Func),
        ("display", prims::display as Func),
        ("newline", prims::newline as Func),
    ] {
        vars.insert(name.to_owned(), Value::Func(fun));
    }

    Env {
        vars,
        enclosing: None,
    }
}

mod prims {
    use crate::{err, RTError, Value};

    // TODO: DRY +-/*
    pub(crate) fn add(args: &[Value]) -> Result<Value, RTError> {
        args.iter()
            .map(Value::as_num)
            .try_fold(0.0, |a, b| Ok(a + b?))
            .map(Value::Num)
    }

    pub(crate) fn mul(args: &[Value]) -> Result<Value, RTError> {
        args.iter()
            .map(Value::as_num)
            .try_fold(1.0, |a, b| Ok(a * b?))
            .map(Value::Num)
    }

    pub(crate) fn div(args: &[Value]) -> Result<Value, RTError> {
        let init = args
            .get(0)
            .ok_or_else(|| RTError("`div` needs at least one argument".to_owned()))?
            .as_num()?;

        Ok(Value::Num(if args.len() == 1 {
            1.0 / init
        } else {
            let rest = mul(&args[1..])?.as_num().unwrap();
            init / rest
        }))
    }

    pub(crate) fn sub(args: &[Value]) -> Result<Value, RTError> {
        let init = args
            .get(0)
            .ok_or_else(|| RTError("`sub` needs at least one argument".to_owned()))?
            .as_num()?;

        Ok(Value::Num(if args.len() == 1 {
            -init
        } else {
            let rest = add(&args[1..])?.as_num().unwrap();
            init - rest
        }))
    }

    pub(crate) fn equals(args: &[Value]) -> Result<Value, RTError> {
        Ok(Value::Bool(args.array_windows().all(|[l, r]| l == r)))
    }

    pub(crate) fn display(args: &[Value]) -> Result<Value, RTError> {
        let [arg] = args else {return err("To many args to `display`".to_owned())};
        print!("{:?}", arg);
        Ok(Value::NotAValue)
    }

    pub(crate) fn newline(args: &[Value]) -> Result<Value, RTError> {
        if !args.is_empty() {
            return err("Newline takes no args".to_owned());
        }
        println!("");
        Ok(Value::NotAValue)
    }
}

// #[cfg(test)]
// mod tests {
//     use super::*;

//     #[test]
//     fn simple_math_space() {
//         let t = grammar::TreeParser::new()
//             .parse("( + 1 2 ( / 2 3 4 5) )")
//             .unwrap();
//         assert_eq!(
//             t,
//             Tree::Add(vec![
//                 Tree::Val(1.0),
//                 Tree::Val(2.0),
//                 Tree::Div(vec![
//                     Tree::Val(2.0),
//                     Tree::Val(3.0),
//                     Tree::Val(4.0),
//                     Tree::Val(5.0),
//                 ])
//             ])
//         );
//     }

//     #[test]
//     fn simple_math_dence() {
//         let t = grammar::TreeParser::new()
//             .parse("(+ 1 2 (/ 2 3 4 5))")
//             .unwrap();
//         assert_eq!(
//             t,
//             Tree::Add(vec![
//                 Tree::Val(1.0),
//                 Tree::Val(2.0),
//                 Tree::Div(vec![
//                     Tree::Val(2.0),
//                     Tree::Val(3.0),
//                     Tree::Val(4.0),
//                     Tree::Val(5.0),
//                 ])
//             ])
//         );
//     }
// }

#[derive(Completer, Helper, Highlighter, Hinter)]
struct InputValidator {
    brackets: MatchingBracketValidator,
}

impl Validator for InputValidator {
    fn validate(&self, ctx: &mut ValidationContext) -> rustyline::Result<ValidationResult> {
        self.brackets.validate(ctx)
    }
}