/* relish: versatile lisp shell
 * Copyright (C) 2021 Aidan Hahn
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

use crate::eval::eval;
use crate::error::{Traceback, start_trace};
use crate::segment::{Ctr, Seg, Type};
use std::collections::HashMap;
use std::fmt;
use std::rc::Rc;

#[derive(Clone)]
pub struct SymTable(HashMap<String, Symbol>, usize);

#[derive(Debug, Clone)]
pub struct UserFn {
    // Un-evaluated abstract syntax tree
    pub ast: Box<Seg>,
    // list of argument string tokens
    pub arg_syms: Vec<String>,
}

/* A symbol may either be a pointer to a function
 * or a syntax tree to eval with the arguments or
 * a simple variable declaration (which can also
 * be a macro)
 */
#[derive(Clone)]
pub enum ValueType {
    Internal(Rc<dyn Fn(&Seg, &mut SymTable) -> Result<Ctr, Traceback>>),
    FuncForm(UserFn),
    VarForm(Box<Ctr>),
}

/* Function Args
 * If Lazy, is an integer denoting number of args
 * If Strict, is a list of type tags denoting argument type.
 */
#[derive(Clone)]
pub enum Args {
    Lazy(u128),
    Strict(Vec<Type>),
    Infinite,
    None,
}

#[derive(Clone)]
pub struct Symbol {
    pub value: ValueType,
    pub name: String,
    pub args: Args,
    // for internal control flow constructs
    pub conditional_branches: bool,
    pub docs: String,
    // see SymTable::Insert
    // (only pub begrudgingly)
    pub __generation: usize,
}

impl SymTable {
    pub fn new() -> SymTable {
        SymTable(HashMap::<String, Symbol>::new(), 0)
    }

    pub fn get(&self, arg: &String) -> Option<&Symbol> {
        self.0.get(arg)
    }

    pub fn contains_key(&self, arg: &String) -> bool {
        self.0.contains_key(arg)
    }

    pub fn insert(&mut self, k: String, mut v: Symbol) -> Option<Symbol> {
        self.1 += 1;
        v.__generation = self.1;
        self.0.insert(k, v)
    }

    pub fn remove(&mut self, arg: &String) -> Option<Symbol> {
        self.0.remove(arg)
    }

    pub fn iter(&self) -> std::collections::hash_map::Iter<'_, String, Symbol> {
        self.0.iter()
    }

    pub fn keys(&self) -> std::collections::hash_map::Keys<String, Symbol> {
        self.0.keys()
    }

    pub fn update(&mut self, other: &mut SymTable) {
        /* updates self with all syms in other that match the following cases:
         *   * sym is not in self
         *   * sym has a newer generation than the entry in self
         */
        let tmp = self.1;
        for i in other.iter() {
            self.0.entry(i.0.to_string())
                .and_modify(|inner: &mut Symbol| {
                    if tmp < i.1.__generation {
                        inner.__generation = i.1.__generation;
                        inner.value = i.1.value.clone();
                        inner.args = i.1.args.clone();
                        inner.docs = i.1.docs.clone();
                        inner.conditional_branches = i.1.conditional_branches;
                        inner.name = i.1.name.clone();
                    }
                })
                .or_insert(i.1.clone());
        }
    }

    pub fn call_symbol(
        &mut self,
        name: &String,
        args: &Seg,
        call_func: bool,
    ) -> Result<Box<Ctr>, Traceback> {
        let mut symbol = match self.remove(name) {
            Some(s) => s,
            None => return Err(
                Traceback::new()
                    .with_trace((name, "(is an undefined symbol)").into())
            ),
        };
        // will re-increment when inserted
        // but we dont want to increment it
        symbol.__generation -= 1;
        self.insert(name.to_string(), symbol.clone());
        if let ValueType::VarForm(ref val) = symbol.value {
            match **val {
                Ctr::Lambda(ref l) if call_func => {
                    return call_lambda(
                        l,
                        &Box::new(Ctr::Seg(args.clone())),
                        self
                    )
                },
                Ctr::Symbol(ref s) if self.is_function_type(s).is_some()
                    && self.is_function_type(s).unwrap() => {
                        symbol = match self.remove(s) {
                            Some(sym) => sym,
                            None => return Err(
                                Traceback::new().with_trace(
                                    (name, format!("(references undefined symbol {})", s))
                                        .into())),
                        };
                        symbol.__generation -= 1;
                        self.insert(symbol.name.clone(), symbol.clone());
                    },
                _ => return Ok(val.clone()),
            }
        }
        if call_func {
            symbol.call(args, self)
        } else {
            // its a function but call_func is off
            Ok(Box::new(Ctr::Symbol(name.to_string())))
        }
    }

    pub fn is_function_type(&self, name: &String) -> Option<bool> {
        if let ValueType::VarForm(ref val) = self.get(name)?.value {
            match **val {
                Ctr::Lambda(_) => Some(true),
                Ctr::Symbol(ref n) => self.is_function_type(n),
                _ =>  Some(false),
            }
        } else {
            Some(true)
        }
    }
}

impl Default for SymTable {
    fn default() -> Self {
        Self::new()
    }
}

impl Args {
    fn validate_inputs(&self, args: &Seg, caller: &String) -> Result<(), Traceback> {
        match self {
            Args::None => {
                if args.len() == 1 {
                    if let Ctr::None = *args.car {
                        return Ok(());
                    } else {
                        return Err(start_trace(
                            (caller, "expected no args")
                                .into()))
                    }
                } else {
                    return Err(start_trace(
                        (caller, "expected no args")
                            .into()))
                }
            }

            Args::Infinite => {
                if !args.is_empty() {
                    return Ok(());
                } else {
                    return Err(start_trace(
                            (caller, "expected args, but none were provided")
                                .into()))
                }
            }

            Args::Lazy(ref num) => {
                let called_arg_count = args.len();
                if *num == 0 {
                    if let Ctr::None = *args.car {
                        //pass
                    } else {
                        return Err(start_trace(
                            (caller, "expected no args, got 1 or more")
                                .into()))
                    }
                } else if *num != called_arg_count {
                    return Err(start_trace(
                        (caller, format!("expected {} args. Got {}.", num, called_arg_count))
                            .into()))
                } else if let Ctr::None = *args.car {
                    return Err(start_trace(
                        (caller, format!("expected {} args. Got 0.", num))
                            .into()))
                }
            }

            Args::Strict(ref arg_types) => {
                let mut idx: usize = 0;
                let mut mismatch = false;
                let passes = args.circuit(&mut |c: &Ctr| -> bool {
                    if idx >= arg_types.len() {
                        return false;
                    }
                    if let Ctr::None = c {
                        return false;
                    }
                    if arg_types[idx] == c.to_type() {
                        idx += 1;
                        return true;
                    }
                    mismatch = true;
                    false
                });

                if passes && idx < (arg_types.len() - 1) {
                    return Err(start_trace(
                        (caller, format!("{} too few arguments", arg_types.len() - (idx + 1)))
                            .into()))
                }

                if !passes {
                    if mismatch {
                        return Err(start_trace(
                            (caller, format!("arg {} expected to be {}", idx + 1, arg_types[idx]))
                                .into()))
                    }
                    if idx > (arg_types.len() - 1) {
                        return Err(start_trace(
                            (caller, "too many arguments".to_string())
                                .into()))
                    }
                    if idx < (arg_types.len() - 1) {
                        return Err(start_trace(
                            (caller, "too few arguments".to_string())
                                .into()))
                    }
                }
            }
        }

        Ok(())
    }
}

impl fmt::Display for Args {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Args::None => write!(f, "none"),
            Args::Infinite => write!(f, "infinite, untyped"),
            Args::Lazy(n) => write!(f, "{} args of any type", n),
            Args::Strict(s) => {
                write!(f, "types: ")?;
                for arg in s {
                    write!(f, "{} ", arg)?
                }
                Ok(())
            }
        }
    }
}

impl fmt::Display for UserFn {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "(lambda (")?;
        let mut arg_iter = (&self.arg_syms).into_iter();
        if let Some(elem) = arg_iter.next() {
            write!(f, "{}", elem)?;
            for i in arg_iter {
                write!(f, " {}", i)?;
            }
        }
        write!(f, ") {})", self.ast.car)?;

        Ok(())
    }
}

impl fmt::Display for ValueType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ValueType::VarForm(ref s) => write!(f, "{}", s),
            ValueType::Internal(_) => write!(f, "<builtin>"),
            ValueType::FuncForm(ref form) => {
                write!(f, "args: ")?;
                for sym in form.arg_syms.clone() {
                    write!(f, "{} ", sym)?;
                }
                write!(f, "\nform: {}", form.ast)?;
                Ok(())
            }
        }
    }
}

impl Symbol {
    /* call
     * routine is called by eval when a symbol is expanded
     */
    pub fn call(&self, args: &Seg, syms: &mut SymTable) -> Result<Box<Ctr>, Traceback> {
        let evaluated_args: &Seg;
        let mut outer_scope_seg_storage = Seg::new();
        let mut cursor = 0;
        let mut errcon: Traceback = Traceback::new();
        if !self.conditional_branches {
            if !args.circuit(&mut |arg: &Ctr| -> bool {
                if let Ctr::Seg(ref s) = arg {
                    let eval_res = eval(s, syms);
                    if eval_res.is_err() {
                        errcon = eval_res.err().unwrap();
                        return false
                    }
                    outer_scope_seg_storage.append(eval_res.unwrap());
                } else if let Ctr::Symbol(ref s) = arg {
                    let eval_res = syms.call_symbol(
                        s,
                        &outer_scope_seg_storage,
                        false
                    );
                    if eval_res.is_err() {
                        errcon = eval_res.err().unwrap();
                        return false
                    }
                    outer_scope_seg_storage.append(eval_res.unwrap());
                } else {
                    outer_scope_seg_storage.append(Box::new(arg.clone()));
                }
                cursor += 1;
                true
            }) {
                return Err(
                    errcon.with_trace((
                        &self.name,
                        format!("error evaluating arg {cursor}"),
                    ).into()))
            }
            evaluated_args = &outer_scope_seg_storage;
        } else {
            evaluated_args = args;
        };

        self.args.validate_inputs(evaluated_args, &self.name)?;

        match &self.value {
            ValueType::VarForm(ref f) => Ok(Box::new(*f.clone())),
            ValueType::Internal(ref f) => Ok(Box::new(f(evaluated_args, syms)?)),
            ValueType::FuncForm(ref f) => {
                // stores any value overwritten by local state
                // If this ever becomes ASYNC this will need to
                // become a more traditional stack design, and the
                // global table will need to be released
                 let mut holding_table = SymTable::new();

                // Prep var table for function execution
                for n in 0..f.arg_syms.len() {
                    if let Some(old) = syms.insert(
                        f.arg_syms[n].clone(),
                        Symbol {
                            name: f.arg_syms[n].clone(),
                            value: ValueType::VarForm(Box::new(evaluated_args[n].clone())),
                            args: Args::None,
                            docs: format!("local argument to {}", f.arg_syms[n].clone()),
                            conditional_branches: false,
                            ..Default::default()
                        },
                    ) {
                        holding_table.insert(f.arg_syms[n].clone(), old);
                    }
                }

                // execute function
                let mut result: Box<Ctr>;
                let mut iterate = &*(f.ast);
                loop {
                    if let Ctr::Seg(ref data) = *iterate.car {
                        match eval(data, syms) {
                            Ok(ctr) => result = ctr,
                            Err(e) => return Err(e.with_trace((&self.name, "returned error").into())),
                        }
                    } else {
                        let temp = Seg::from_mono(iterate.car.clone());
                        match eval(&temp, syms) {
                            Ok(ctr) => {
                                if let Ctr::Seg(s) = *ctr {
                                    result = s.car.clone();
                                } else {
                                    result = ctr;
                                }
                            }
                            Err(e) => return Err(
                                e.with_trace((&self.name, "returned error").into())
                            ),
                        }
                    }

                    match *iterate.cdr {
                        Ctr::Seg(ref next) => iterate = next,
                        Ctr::None => break,
                        _ => panic!("function body not in standard form!"),
                    }
                }

                // clear local vars and restore previous values
                for n in 0..f.arg_syms.len() {
                    syms.remove(&f.arg_syms[n]);
                    if let Some(val) = holding_table.remove(&f.arg_syms[n]) {
                        syms.insert(f.arg_syms[n].clone(), val);
                    }
                }

                Ok(result)
            }
        }
    }

    pub fn from_ast(
        id: &String,
        doc: &String,
        ast: &Seg,
        arg_list: Option<Vec<String>>,
    ) -> Symbol {
        let args: Args;
        let value: ValueType;

        if let Some(ref arg_syms) = arg_list {
            value = ValueType::FuncForm(UserFn{
                ast: Box::new(ast.clone()),
                arg_syms: arg_syms.clone(),
            });
            args = Args::Lazy(arg_syms.len() as u128);
        } else {
            args = Args::None;
            value = ValueType::VarForm(ast.car.clone());
        }

        Symbol {
            name: id.clone(),
            docs: doc.clone(),
            conditional_branches: false,
            args, value,
            ..Default::default()
        }
    }
}

impl Default for Symbol {
    fn default() -> Symbol {
        Symbol {
            value: ValueType::Internal(
                Rc::new(
                    |_: &Seg, _: &mut SymTable| -> Result<Ctr, Traceback> {
                        unimplemented!()
                    }
                )
            ),
            name: String::new(),
            docs: String::new(),
            args: Args::None,
            conditional_branches: false,
            __generation: 0,
        }
    }
}

pub fn call_lambda(
    lam: &UserFn,
    args_ctr: &Box<Ctr>,
    syms: &mut SymTable,
) -> Result<Box<Ctr>, Traceback> {
    let temp_sym = Symbol {
        name: String::from("<lambda>"),
        conditional_branches: false,
        docs: String::from("user defined lambda"),
        args: Args::Lazy(lam.arg_syms.len() as u128),
        value: ValueType::FuncForm(lam.clone()),
        ..Default::default()
    };

    let args: &Seg;
    let outer_scope_maybe_args: Seg;
    if let Ctr::Seg(ref args_head) = **args_ctr {
        args = args_head;
    } else {
        outer_scope_maybe_args = Seg::from_mono(
            Box::new(*args_ctr.clone()));
        args = &outer_scope_maybe_args;
    }

    temp_sym.call(args, syms)
}