openbirch-rs/src/cli-repl.rs

use std::fs::File;
use std::io::{self, Read, StdoutLock, Write, stdout};
use std::rc::Rc;
use std::{env, f64};

use libopenbirch::environment::{Environment, EnvironmentInternalSymbolKey};
use libopenbirch::node::call::Call;
use libopenbirch::node::constant::Constant;
use libopenbirch::node::empty::Empty;
use libopenbirch::node::set::Set;
use libopenbirch::node::string_node::StringNode;
use libopenbirch::node::{Error, Node, NodeEnum};
use libopenbirch::parser::{Lexer, LexerError, Parser, ParserError};
use termion::color;
use termion::event::Key;
use termion::input::TermReadEventsAndRaw;
use termion::raw::{IntoRawMode, RawTerminal};

pub struct Input {
    #[cfg(not(feature = "async"))]
    stdin: std::io::Stdin,

    stdout: RawTerminal<StdoutLock<'static>>,
    buffer: String,
    current_char: usize,

    history: Vec<String>,
    history_idx: Option<usize>,
}
impl Input {
    pub fn new() -> Self {
        #[cfg(not(feature = "async"))]
        let stdin = std::io::stdin();
        let stdout = stdout().lock().into_raw_mode().unwrap();

        Self {
            stdin,
            stdout,
            buffer: "".into(),
            current_char: 0,

            history: vec![],
            history_idx: None,
        }
    }

    #[inline]
    fn draw_line(
        buffer: &String,
        stdout: &mut RawTerminal<StdoutLock<'static>>,
        current_char: usize,
    ) {
        print!("\r{}> {}", termion::clear::CurrentLine, buffer,);
        let left_diff = buffer.len() - current_char;
        if left_diff > 0 {
            print!("{}", termion::cursor::Left(left_diff.try_into().unwrap()));
        }
        let _ = stdout.flush();
    }

    /// Gets input from `io::stdin`.
    /// Returns `None` when the user presses <C-c> or types `:quit`
    pub fn get(&mut self) -> Result<Option<String>, io::Error> {
        Self::draw_line(&self.buffer, &mut self.stdout, self.current_char);
        loop {
            for b in (&mut self.stdin).events_and_raw() {
                let (event, _slice) = b?;

                match event {
                    termion::event::Event::Key(key) => match key {
                        Key::Char(char) if char == '\n' => {
                            if &self.buffer == ":quit" {
                                println!("\r");
                                return Err(io::Error::from(io::ErrorKind::Interrupted));
                            }

                            if self.buffer.trim().len() > 0 {
                                self.history.insert(0, self.buffer.clone());
                            }
                            let r = self.buffer.clone();
                            self.buffer.clear();
                            self.current_char = 0;
                            self.history_idx = None;
                            print!("\n\r");
                            return Ok(Some(r));
                        }
                        Key::Backspace => {
                            if self.current_char == 0 {
                                continue;
                            }
                            self.current_char -= 1;
                            self.buffer.remove(self.current_char);
                            Self::draw_line(&self.buffer, &mut self.stdout, self.current_char);
                            let _ = self.stdout.flush();
                        }
                        Key::Delete => {
                            if self.current_char == self.buffer.len() {
                                continue;
                            }
                            self.buffer.remove(self.current_char);
                            Self::draw_line(&self.buffer, &mut self.stdout, self.current_char);
                            let _ = self.stdout.flush();
                        }
                        Key::Left => {
                            if self.current_char == 0 {
                                continue;
                            }
                            self.current_char -= 1;
                            print!("{}", termion::cursor::Left(1));
                            let _ = self.stdout.flush();
                        }
                        Key::Right => {
                            if self.current_char == self.buffer.len() {
                                continue;
                            }
                            self.current_char += 1;
                            print!("{}", termion::cursor::Right(1));
                            let _ = self.stdout.flush();
                        }
                        Key::Up => {
                            match self.history_idx {
                                Some(i) => self.history_idx = Some((i + 1) % self.history.len()),
                                None => {
                                    if self.history.len() > 0 {
                                        self.history_idx = Some(0)
                                    } else {
                                        break;
                                    }
                                }
                            }

                            self.buffer =
                                self.history.get(self.history_idx.unwrap()).unwrap().clone();
                            self.current_char = self.buffer.len();
                            Self::draw_line(&self.buffer, &mut self.stdout, self.current_char);
                        }
                        Key::Down => {
                            match self.history_idx {
                                Some(i) => {
                                    let i = if i == 0 { self.history.len() } else { i };
                                    self.history_idx = Some((i - 1) % self.history.len())
                                }
                                None => {
                                    if self.history.len() > 0 {
                                        self.history_idx = Some(self.history.len() - 1);
                                    } else {
                                        break;
                                    }
                                }
                            }

                            self.buffer =
                                self.history.get(self.history_idx.unwrap()).unwrap().clone();
                            self.current_char = self.buffer.len();
                            Self::draw_line(&self.buffer, &mut self.stdout, self.current_char);
                        }
                        Key::Home => self.current_char = 0,
                        Key::End => self.current_char = self.buffer.len(),
                        Key::Char(char) => {
                            self.buffer.insert(self.current_char, char);
                            self.current_char += 1;
                            Self::draw_line(&self.buffer, &mut self.stdout, self.current_char);
                        }
                        Key::Ctrl(c) if c == 'c' => {
                            println!("\r");
                            return Err(io::Error::from(io::ErrorKind::Interrupted));
                        }
                        _ => {
                            #[cfg(debug_assertions)]
                            return Err(io::Error::other(format!(
                                "Key {key:?} is not implemented"
                            )));
                        }
                    },
                    termion::event::Event::Mouse(_mouse_event) => {}
                    termion::event::Event::Unsupported(_items) => {}
                }
            }
        }
    }

    pub fn disable_raw(&mut self) -> io::Result<()> {
        self.stdout.suspend_raw_mode()
    }

    pub fn enable_raw(&mut self) -> io::Result<()> {
        self.stdout.activate_raw_mode()
    }
}

fn print_err(i: usize, len: usize, exp: String) {
    println!(
        "\r{}{}{}",
        " ".repeat(i + 3 - len),
        color::Fg(color::Yellow),
        "^".repeat(len)
    );
    println!(
        "\r{}{}{}",
        color::Fg(color::Red),
        exp,
        color::Fg(color::Reset)
    );
}

fn print(args: &Vec<Rc<NodeEnum>>, env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    for expr in args {
        println!("\r{}", expr.as_string(Some(env)));
    }

    Ok(Empty::new(""))
}

fn unset(args: &Vec<Rc<NodeEnum>>, env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    for arg in args {
        if let NodeEnum::StringNode(string) = arg.as_ref() {
            if let Some(id) = env.str_to_id(string.get_value()) {
                env.undefine(*id);
            }
        } else {
            return Err(Error::MismatchedType(format!(
                "Expected strings, but one of the arguments was a {}",
                arg.type_str()
            )));
        }
    }

    Ok(Empty::new(format!("Undefined {} symbols", args.len())))
}

fn get_float_precision(
    args: &Vec<Rc<NodeEnum>>,
    env: &mut Environment,
) -> Result<Rc<NodeEnum>, Error> {
    match args.len() {
        0 => Ok(Rc::new(
            Constant::new_from_float(env.get_float_precision() as f64, &env).into(),
        )),
        1 => {
            let arg = args.first().unwrap();
            if let NodeEnum::Constant(constant) = arg.as_ref() {
                let prec = constant.get_value().prec_64() as f64;

                Ok(Rc::new(Constant::new_from_float(prec, &env).into()))
            } else {
                Err(Error::MismatchedType(format!(
                    "Expected argument of type Constant, but got {}",
                    arg.type_str()
                )))
            }
        }
        _ => Err(Error::Other(format!(
            "Expected 0 or 1 arguments, got {}",
            args.len()
        )))?,
    }
}

fn set_float_precision(
    args: &Vec<Rc<NodeEnum>>,
    env: &mut Environment,
) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 1 {
        Err(Error::ArgumentCount(1, args.len()))?
    }

    let arg = args.first().unwrap();

    let precision: u64 = if let NodeEnum::Constant(value) = arg.as_ref() {
        value.get_value().to_u32_saturating().unwrap().into()
    } else {
        Err(Error::MismatchedType(format!(
            "Expected argument of type Constant, but got {}",
            arg.type_str()
        )))?
    };

    env.set_float_precision(precision);

    Ok(Empty::new(format!("Set float precision to {precision}")))
}

fn map(args: &Vec<Rc<NodeEnum>>, env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 2 {
        Err(Error::ArgumentCount(2, args.len()))?
    }

    let arg = args.first().unwrap();
    let func = match arg.as_ref() {
        NodeEnum::Function(_) | NodeEnum::Closure(_) => arg,
        _ => Err(Error::MismatchedType(format!(
            "Argument 1 expected a Function but got {}",
            arg.type_str()
        )))?,
    };

    let arg = args.get(1).unwrap().as_ref();
    let set = if let NodeEnum::Set(set) = arg {
        set.get_values()
    } else {
        return Err(Error::MismatchedType(format!(
            "Argument 1 expected a Set but got {}",
            arg.type_str()
        )))?;
    };

    let mut out = Vec::with_capacity(set.len());

    for expr in set.iter() {
        out.push(Call::new(func.clone(), vec![expr.clone()]).evaluate(env)?);
    }

    Ok(Set::new(out))
}

fn get(args: &Vec<Rc<NodeEnum>>, _env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 2 {
        Err(Error::ArgumentCount(2, args.len()))?
    }

    let set = if let NodeEnum::Set(set) = args.first().unwrap().as_ref() {
        set.get_values()
    } else {
        return Err(Error::MismatchedType(format!(
            "Argument 1 expected Set but got {}",
            args.first().unwrap().type_str()
        )));
    };

    let idx = if let NodeEnum::Constant(c) = args.get(1).unwrap().as_ref() {
        if let Some(u) = c.get_value().to_u32_saturating() {
            u as usize
        } else {
            return Err(Error::ConversionError(
                "Error occured while trying to convert second argument to integer".to_owned(),
            ));
        }
    } else {
        return Err(Error::MismatchedType(
            "Argument 1 is expected to be a Set".to_owned(),
        ));
    };

    if let Some(v) = set.get(idx) {
        return Ok(v.clone());
    }

    return Err(Error::IndexOutOfBounds(set.len(), idx));
}

fn length(args: &Vec<Rc<NodeEnum>>, env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 1 {
        Err(Error::ArgumentCount(1, args.len()))?
    }

    let set = if let NodeEnum::Set(set) = args.first().unwrap().as_ref() {
        set.get_values()
    } else {
        return Err(Error::MismatchedType(
            "Argument 1 is expected to be a Set".to_owned(),
        ));
    };

    Ok(Rc::new(
        Constant::new_from_float(set.len() as f32, &env).into(),
    ))
}

fn benchmark(args: &Vec<Rc<NodeEnum>>, env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 1 {
        Err(Error::ArgumentCount(1, args.len()))?
    }

    use std::time::Instant;
    let now = Instant::now();

    let result = Call::new(args.first().unwrap().clone(), vec![]).evaluate(env)?;

    let elapsed = now.elapsed();

    let time_str = StringNode::new(elapsed.as_secs_f64().to_string()).into();
    let time_const = Rc::new(Constant::new_from_float(elapsed.as_secs_f64(), &env).into());

    Ok(Set::new(vec![time_str, time_const, result]))
}

fn head(args: &Vec<Rc<NodeEnum>>, _env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 1 {
        Err(Error::ArgumentCount(1, args.len()))?
    }

    let set = if let NodeEnum::Set(set) = args.first().unwrap().as_ref() {
        set.get_values()
    } else {
        return Err(Error::MismatchedType(format!(
            "Expected a Set but got a {}",
            args.first().unwrap().type_str()
        )));
    };

    let head = set.first().cloned().unwrap_or(Set::new(vec![]));
    let len = set.len();
    let tail = if len > 1 {
        set[1..len].to_vec()
    } else {
        vec![]
    };

    Ok(Set::new(vec![head, Set::new(tail)]))
}

fn mult_equal(args: &Vec<Rc<NodeEnum>>, env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    let a1 = if let NodeEnum::Multiply(m) = args.get(0).ok_or(Error::ArgumentCount(1, 0))?.as_ref()
    {
        m
    } else {
        return Err(Error::MismatchedType(
            "Argument 1 is not a multiply".to_owned(),
        ));
    };
    let a2 = if let NodeEnum::Multiply(m) = args.get(1).ok_or(Error::ArgumentCount(2, 1))?.as_ref()
    {
        m
    } else {
        return Err(Error::MismatchedType(
            "Argument 2 is not a multiply".to_owned(),
        ));
    };

    Ok(Rc::new(NodeEnum::Bool(a1.symbols_match(a2).into())))
}

fn diff(args: &Vec<Rc<NodeEnum>>, _env: &mut Environment) -> Result<Rc<NodeEnum>, Error> {
    if args.len() != 2 {
        return Err(Error::ArgumentCount(2, args.len()));
    }

    let expr = args.first().unwrap().clone();

    let var = if let NodeEnum::Symbol(symbol) = args.get(1).unwrap().as_ref() {
        symbol.get_value()
    } else {
        return Err(Error::MismatchedType(format!(
            "Argument 2 is expected to be a Symbol but got a {}",
            args.get(1).unwrap().type_str()
        )));
    };

    todo!()
}

fn diff_internally(
    expr: Rc<NodeEnum>,
    var: EnvironmentInternalSymbolKey,
    env: &mut Environment,
) -> Result<Rc<NodeEnum>, String> {
    todo!()
}

fn evaluate(source: &String, mut env: &mut Environment, verbose: bool) {
    let mut lexer = Lexer::new(&source);
    let tokens_result = lexer.lex();

    if tokens_result.is_err() {
        match tokens_result.err().unwrap() {
            LexerError::UnexpectedChar(i, exp) => print_err(i, 1, exp),
        }
        return;
    }

    let tokens = tokens_result.unwrap();
    let mut parser = Parser::new(tokens, &mut env);

    let nodes = match parser.parse() {
        Ok(nodes) => nodes,
        Err(err) => {
            match err {
                ParserError::UnexpectedEndOfTokens(exp) => print_err(source.len(), 1, exp),
                ParserError::UnexpectedToken(i, len, exp)
                | ParserError::Unimplemented(i, len, exp) => print_err(i, len, exp),
                ParserError::UnexpectedNode(i, exp) | ParserError::NumberParse(i, exp) => {
                    print_err(i, 1, exp)
                }
            }
            return;
        }
    };

    print!("{}", color::Fg(color::Blue));

    for node in nodes {
        let evaluated = node.evaluate(&mut env);
        match evaluated {
            Ok(result) if verbose => println!("\r\t{}", result.as_string(Some(&env))),
            Err(exp) => print_err(0, 1, exp.to_string()),
            _ => {}
        }
    }
}

fn main() -> Result<(), io::Error> {
    let mut env = Environment::new();

    env.set_float_precision(128);

    env.define_native_function("print", print);
    env.define_native_function("unset", unset);
    env.define_native_function("get_float_precision", get_float_precision);
    env.define_native_function("set_float_precision", set_float_precision);
    env.define_native_function("map", map);
    env.define_native_function("head", head);
    env.define_native_function("get", get);
    env.define_native_function("length", length);
    env.define_native_function("benchmark", benchmark);
    env.define_native_function("mult_equals", mult_equal);

    env.define(
        "pi",
        Rc::new(Constant::new_from_float(std::f64::consts::PI, &env).into()),
    );

    let args: Vec<String> = std::env::args().collect();

    if args.len() > 1 {
        let file = args.get(1).unwrap();

        let mut file = File::open(file)?;
        let mut src = String::new();
        let _ = file.read_to_string(&mut src);

        evaluate(&src, &mut env, false);

        return Ok(());
    }

    let mut input = Input::new();

    while let Some(source) = input.get()? {
        input.disable_raw()?;

        evaluate(&source, &mut env, true);

        input.enable_raw()?;

        print!("{}", color::Fg(color::Reset));

        // println!("\t{}{source}{}", termion::color::Fg(termion::color::Blue), termion::color::Fg(termion::color::Reset));
    }

    Ok(())
}