hence/src/lib/arg.rs

use anyhow::{bail, Result};

use crate::arg;
use crate::assembler;
use crate::lexer;

#[derive(Debug, Clone)]
pub enum Arg {
    String(String),
    Number(u16),
    Variable(String),
    BinaryExpression {
        left: Box<Self>,
        right: Box<Self>,
        op: BinaryExpressionOperator,
    },
}

impl assembler::ToAssembly for Arg {
    fn to_assembly(&self) -> String {
        match self {
            Arg::String(x) => format!("\"{x}\""),
            Arg::Number(x) => x.to_string(),
            Arg::Variable(x) => x.clone(),
            Arg::BinaryExpression { left, right, op } => {
                format!(
                    "({left} {op} {right})",
                    left = left.to_assembly(),
                    op = op.to_assembly(),
                    right = right.to_assembly()
                )
            }
        }
    }
}

impl assembler::ByteResolvable<assembler::Data> for Arg {
    fn resolve_number(&self, data: &mut assembler::Data) -> Result<u16, String> {
        match self {
            Arg::String(x) => {
                let y = x.replace("\\n", "\n");

                if y.len() == 1 {
                    Ok(y.as_bytes()[0] as u16)
                } else {
                    let bytes = y.as_bytes();

                    Ok(((bytes[0] as u16) << 8) | bytes[1] as u16)
                }
            }
            Arg::Number(x) => Ok(*x),
            Arg::Variable(x) => {
                let mut name = x.clone();
                let mut arg: Option<arg::Arg>;

                loop {
                    arg = data.constants.get(&name).cloned();

                    match arg {
                        Some(a) => {
                            // arg = Some(a.clone());
                            match a {
                                arg::Arg::Variable(n) => {
                                    name = n;
                                }
                                _ => break Ok(a.resolve_number(data).unwrap()),
                            };
                        }
                        _ => return Err(format!("Variable does not exist: '{name}'")),
                    }
                }
            }
            Arg::BinaryExpression { left, right, op } => {
                let left_val = left.resolve_number(data).unwrap();
                let right_val = right.resolve_number(data).unwrap();

                Ok(match op {
                    BinaryExpressionOperator::Add => left_val.wrapping_add(right_val),
                    BinaryExpressionOperator::Sub => left_val.wrapping_sub(right_val),
                    BinaryExpressionOperator::Mul => left_val.wrapping_mul(right_val),
                    BinaryExpressionOperator::Div => left_val.wrapping_div(right_val),
                    BinaryExpressionOperator::Pow => left_val.wrapping_pow(right_val as u32),
                })
            }
        }
    }

    fn resolve_bytes(&self, data: &mut assembler::Data) -> Result<Vec<u8>, String> {
        match self {
            Arg::String(x) => Ok(x.replace("\\n", "\n").bytes().collect()),
            Arg::Number(x) => Ok(vec![(*x >> 8) as u8, *x as u8]),
            Arg::Variable(x) => {
                let mut name = x.clone();
                let mut arg: Option<arg::Arg>;

                loop {
                    arg = data.constants.get(&name).cloned();

                    match arg {
                        Some(a) => {
                            // arg = Some(a.clone());
                            match a {
                                arg::Arg::Variable(n) => {
                                    name = n;
                                }
                                _ => break Ok(a.resolve_bytes(data).unwrap()),
                            };
                        }
                        _ => return Err(format!("Variable does not exist: '{name}'")),
                    }
                }
            }
            Arg::BinaryExpression {
                left: _,
                right: _,
                op: _,
            } => {
                let x = self.resolve_number(data).unwrap();

                Ok(vec![(x >> 8) as u8, x as u8])
            }
        }
    }
}

#[derive(Debug, Clone)]
pub enum BinaryExpressionOperator {
    Add,
    Sub,
    Mul,
    Div,
    Pow,
}

impl assembler::ToAssembly for BinaryExpressionOperator {
    fn to_assembly(&self) -> String {
        match self {
            BinaryExpressionOperator::Add => "+".to_string(),
            BinaryExpressionOperator::Sub => "-".to_string(),
            BinaryExpressionOperator::Mul => "*".to_string(),
            BinaryExpressionOperator::Div => "/".to_string(),
            BinaryExpressionOperator::Pow => "**".to_string(),
        }
    }
}

pub fn parse_binary_operation(token: &lexer::Token) -> Result<BinaryExpressionOperator> {
    match token {
        lexer::Token::Add => Ok(BinaryExpressionOperator::Add),
        lexer::Token::Sub => Ok(BinaryExpressionOperator::Sub),
        lexer::Token::Mul => Ok(BinaryExpressionOperator::Mul),
        lexer::Token::Div => Ok(BinaryExpressionOperator::Div),
        lexer::Token::Pow => Ok(BinaryExpressionOperator::Pow),
        _ => bail!("Invalid binary expression operator"),
    }
}

pub fn parse_binary_expression_arg(tokens: &mut Vec<&&lexer::Token>) -> Result<Arg> {
    if tokens.is_empty() {
        bail!("Malformed binary expression");
    }

    let mut args: Vec<&&lexer::Token> = tokens.drain(..3).collect();

    let mut bin_expr = Arg::BinaryExpression {
        left: Box::new((&parse_args(vec![args[0]]).unwrap()[0]).clone()),
        right: Box::new((&parse_args(vec![args[2]]).unwrap()[0]).clone()),
        op: parse_binary_operation(args[1]).unwrap(),
    };

    while !tokens.is_empty() {
        args = tokens.drain(..2).collect();
        bin_expr = Arg::BinaryExpression {
            left: Box::new(bin_expr),
            right: Box::new((&parse_args(vec![args[1]]).unwrap()[0]).clone()),
            op: parse_binary_operation(args[0]).unwrap(),
        }
    }

    Ok(bin_expr)
}

pub fn parse_args(tokens: Vec<&lexer::Token>) -> Result<Vec<Arg>, &str> {
    let mut iter = tokens.iter();
    let mut args: Vec<Arg> = Vec::new();
    let mut last_was_comma = true;

    while let Some(token) = iter.next() {
        match token {
            lexer::Token::Comma => {
                last_was_comma = true;
            }
            _ => {
                if !last_was_comma {
                    return Err("Invalid argument separation");
                }

                match token {
                    lexer::Token::Comment(_) => {}
                    lexer::Token::StringLiteral(x) => {
                        args.push(Arg::String(x.clone()));
                    }
                    lexer::Token::Literal(x) => {
                        args.push(Arg::Variable(x.clone()));
                    }
                    lexer::Token::Number(x) => {
                        args.push(Arg::Number(match num_parse::parse_uint(x) {
                            Some(y) => y,
                            _ => return Err("Error parsing number"),
                        }));
                    }
                    lexer::Token::LParen => {
                        let mut depth: usize = 1;

                        args.push(
                            parse_binary_expression_arg(
                                &mut iter
                                    .by_ref()
                                    .take_while(|t| match t {
                                        lexer::Token::LParen => {
                                            depth += 1;
                                            true
                                        }
                                        lexer::Token::RParen => {
                                            depth -= 1;
                                            depth != 0
                                        }
                                        _ => true,
                                    })
                                    .collect(),
                            )
                            .unwrap(),
                        );
                    }
                    _ => return Err("Unexpected token for argument"),
                }

                last_was_comma = false;
            }
        }
    }

    Ok(args)
}