glowdust 0.0.1

use crate::compiler::parser::ast::{BinOpToken, ExprKind, LiteralToken, Span, Stmt, UnaryOpToken};
use crate::compiler::parser::ast::{Expr, FunctionDefinition};
use crate::compiler::parser::parse_state::{FunctionContext, ParseState};
use crate::compiler::parser::pointer::P;
use crate::compiler::query_solver::solve_query;
use crate::compiler::value::parameters::Parameter;
use crate::compiler::value::values::{FunctionDefinitionObject, TypeDefinition, Value};
use crate::runtime::bytecode::OpCode::{
    OpAdd, OpAnd, OpCall, OpConst, OpDivide, OpEquals, OpFalse, OpGetField, OpGetGlobal,
    OpGetLocal, OpGt, OpLt, OpMktuple, OpMultiply, OpNegate, OpNewInstance, OpNot, OpOr, OpPop,
    OpProduce, OpReturn, OpSetField, OpSetFunt, OpSetGlobal, OpSetLocal, OpSubtract, OpTrue,
    OpTupleGet,
};
use crate::runtime::compiled_script::{Bytecode, Executable};
use crate::store::Store;
use crate::{boolean_to_value, string_to_value};
use std::collections::HashSet;

use crate::compiler::parser::CompileError;

pub struct WalkContext<'walk_time> {
    pub is_local: bool,
    pub constants: &'walk_time [Value],
    pub initialized_variables: HashSet<u8>,
    pub store: &'walk_time mut dyn Store,
}

fn type_into_fdo(type_def: &TypeDefinition) -> Result<FunctionDefinitionObject, CompileError> {
    let mut synthetic_def = FunctionDefinitionObject::default();
    synthetic_def.name = type_def.name.clone();

    synthetic_def
        .constants
        .push(string_to_value!(synthetic_def.name.clone()));
    synthetic_def.bytecode.byte(OpNewInstance as u8);
    synthetic_def.bytecode.byte(0); // the name of the type is constant 0, we added it not two lines ago

    for (count, (field_name, _)) in type_def.fields.iter().enumerate() {
        synthetic_def
            .formal_arguments
            .push(Parameter::Variable(field_name.clone()));
        synthetic_def
            .constants
            .push(string_to_value!(field_name.clone()));
        synthetic_def.bytecode.byte(OpGetLocal as u8);
        synthetic_def.bytecode.byte((count) as u8);
        synthetic_def.bytecode.byte(OpSetField as u8);
        synthetic_def.bytecode.byte((count + 1) as u8);
    }

    synthetic_def.bytecode.byte(OpReturn as u8);

    Ok(synthetic_def)
}

// TODO Obviously this needs cleanup, to make walker a trait and factor out all boilerplate. As they say,
// TODO if I had more time, I would have made this code shorter.
// TODO this would be a good first issue but it needs to be fixed soon

pub fn walk(state: &ParseState, store: &mut dyn Store) -> Result<Executable, CompileError> {
    let mut result = Executable::default();

    for (name, defs) in &state.function_definitions {
        let mut out_defs = vec![];
        for def in defs {
            let fdo = walk_fun_def(&def.0, &def.1, store)?;
            out_defs.push(fdo);
        }
        result.function_definitions.insert(name.clone(), out_defs);
    }

    for type_def in &state.type_definitions {
        let mut current = TypeDefinition::default();

        current.name = type_def.name.clone();
        for field in &type_def.fields {
            current
                .fields
                .push((field.name.clone(), field.type_name.name.clone()));
        }

        let type_constructor_fdo = type_into_fdo(&current)?;
        result
            .function_definitions
            .insert(current.name.clone(), vec![type_constructor_fdo]);
        result.type_definitions.push(current);
    }

    let mut context = WalkContext {
        is_local: false,
        constants: &state.script_symbols.constants,
        initialized_variables: HashSet::new(),
        store,
    };
    // Existing state contains global variables. These need to be marked as initialized.
    for var in &state.script_symbols.variables {
        if var.initialized {
            context.initialized_variables.insert(var.index as u8);
        }
    }
    for stmt in &state.statements {
        validate_variable_initialization_in_stmt(&mut context, stmt)?;
    }
    let mut bites = Bytecode::new();
    for stmt in &state.statements {
        generate_bytecode_from_statement(&mut context, &mut bites, stmt)?;
    }
    bites.byte(OpReturn as u8);
    result.script_bytecode = bites;
    result.script_constants = state.script_symbols.constants.clone();
    result.script_variables = state.script_symbols.variables.clone();
    for var in result.script_variables.iter_mut() {
        if context.initialized_variables.contains(&(var.index as u8)) {
            var.initialized = true;
        }
    }
    Ok(result)
}

pub(crate) fn validate_variable_initialization_in_stmt(
    walk_context: &mut WalkContext,
    statement: &Stmt,
) -> Result<(), CompileError> {
    match statement {
        Stmt::Expr(expr) | Stmt::Semi(expr) => validate_variable_initialization(walk_context, expr),
        Stmt::Return(exprs) => {
            for expr in exprs {
                validate_variable_initialization(walk_context, expr)?
            }
            Ok(())
        }
        Stmt::Match(_) => Ok(() /*This doesn't feel right, but will do for now*/),
        Stmt::BangBang => Ok(()),
        Stmt::Empty => Ok(()),
    }
}

pub(crate) fn validate_variable_initialization(
    walk_context: &mut WalkContext,
    expr: &Expr,
) -> Result<(), CompileError> {
    match &expr.kind {
        ExprKind::Tuple(exprs) => {
            for expr in exprs {
                validate_variable_initialization(walk_context, expr)?;
            }
        }
        ExprKind::Unary(_, expr) => {
            validate_variable_initialization(walk_context, expr)?;
        }
        ExprKind::Binary(_, lhs, rhs) => {
            validate_variable_initialization(walk_context, lhs)?;
            validate_variable_initialization(walk_context, rhs)?;
        }
        ExprKind::Assignment(lhs, rhs) => {
            validate_variable_initialization(walk_context, rhs)?;
            if let ExprKind::Variable(id) = &lhs.kind {
                walk_context.initialized_variables.insert(*id);
            }
        }
        ExprKind::Variable(id) => {
            // this is the first time we see this variable outside of assignment. So it either
            // needs to have been assigned OR it's an inherited global. Either way, it needs to be
            // in the initialized set.
            if !walk_context.initialized_variables.contains(id) {
                return Err(CompileError::new(
                    format!("Variable {} is not initialized yet", *id),
                    expr.span,
                ));
            }
        }
        ExprKind::FunCall(_, params) => {
            for param in params {
                validate_variable_initialization(walk_context, &param.value)?;
            }
        }
        _ => {}
    }
    Ok(())
}

fn walk_fun_def(
    def: &FunctionDefinition,
    symbols: &FunctionContext,
    store: &mut dyn Store,
) -> Result<FunctionDefinitionObject, CompileError> {
    let mut body_bites = Bytecode::new();
    let mut fdo = FunctionDefinitionObject::new(def.name.clone());

    for arg in &def.formal_arguments {
        match &arg.kind {
            ExprKind::Literal(lit) => match lit {
                LiteralToken::True => {
                    fdo.formal_arguments
                        .push(Parameter::Constant(boolean_to_value!(true)));
                }
                LiteralToken::False => {
                    fdo.formal_arguments
                        .push(Parameter::Constant(boolean_to_value!(false)));
                }
                LiteralToken::Integer(id)
                | LiteralToken::Float(id)
                | LiteralToken::StringLiteral(id) => {
                    fdo.formal_arguments.push(Parameter::Constant(
                        symbols.symbols.constants[*id as usize].clone(),
                    ));
                }
            },
            ExprKind::Variable(id) => {
                fdo.formal_arguments.push(Parameter::Variable(
                    symbols
                        .symbols
                        .variables
                        .get(*id as usize)
                        .unwrap()
                        .name
                        .clone(),
                ));
            }
            _ => {
                panic!("Cannot handle {:?} as a function parameter", arg);
            }
        }
    }

    if def.body.is_empty() {
        body_bites.byte_and_line(OpMktuple as u8, 100);
        body_bites.byte(0);
    } else {
        let mut context = WalkContext {
            is_local: true,
            constants: &symbols.symbols.constants,
            initialized_variables: HashSet::new(),
            store,
        };
        for var_param in &def.formal_arguments {
            match var_param.kind {
                ExprKind::Variable(id) => {
                    context.initialized_variables.insert(id);
                }
                _ => { /* Variables in expressions are not used in the body so we only need to worry about variable arguments */
                }
            }
        }
        for stmt in &def.body {
            validate_variable_initialization_in_stmt(&mut context, stmt)?;
        }
        for stmt in &def.body {
            generate_bytecode_from_statement(&mut context, &mut body_bites, stmt)?;
        }
    }
    fdo.bytecode = body_bites;
    fdo.bytecode.byte_and_line(OpReturn as u8, 100);
    fdo.constants = symbols.symbols.constants.clone();
    Ok(fdo)
}

pub(crate) fn generate_bytecode_from_statement(
    walk_context: &mut WalkContext,
    bytecode: &mut Bytecode,
    statement: &Stmt,
) -> Result<(), CompileError> {
    match statement {
        Stmt::Expr(expr) => generate_bytecode(walk_context, bytecode, expr),
        Stmt::Semi(expr) => {
            generate_bytecode(walk_context, bytecode, expr)?;
            bytecode.byte(OpPop as u8);
            Ok(())
        }
        Stmt::Return(values) => {
            for expr in values {
                generate_bytecode(walk_context, bytecode, expr)?;
            }
            bytecode.byte(OpProduce as u8);
            bytecode.byte(values.len() as u8);
            Ok(())
        }
        Stmt::Match(atoms) => {
            let prev = walk_context.is_local;
            walk_context.is_local = true;
            let mut bytes = solve_query(atoms, walk_context)?;
            bytecode.bytes.append(&mut bytes.bytes);
            bytecode.lines.append(&mut bytes.lines);
            walk_context.is_local = prev;
            Ok(())
        }
        Stmt::BangBang => {
            bytecode.byte(OpTrue as u8);
            Ok(())
        }
        Stmt::Empty => Ok(()),
    }
}

pub(crate) fn generate_bytecode(
    walk_context: &mut WalkContext,
    bytecode: &mut Bytecode,
    expr: &Expr,
) -> Result<(), CompileError> {
    match &expr.kind {
        ExprKind::Literal(literal) => match literal {
            LiteralToken::True => {
                bytecode.byte_and_line(OpTrue as u8, expr.span.line);
            }
            LiteralToken::False => {
                bytecode.byte_and_line(OpFalse as u8, expr.span.line);
            }
            LiteralToken::Integer(index)
            | LiteralToken::Float(index)
            | LiteralToken::StringLiteral(index) => {
                bytecode.byte_and_line(OpConst as u8, expr.span.line);
                bytecode.byte_and_line(*index, expr.span.line);
            }
        },
        ExprKind::Binary(op, lhs, rhs) => {
            generate_bytecode(walk_context, bytecode, lhs)?;
            generate_bytecode(walk_context, bytecode, rhs)?;
            op.bytecode(bytecode, expr.span.line);
        }
        ExprKind::Unary(op, expr) => {
            generate_bytecode(walk_context, bytecode, expr)?;
            op.bytecode(bytecode, expr.span.line);
        }
        ExprKind::Assignment(lhs, rhs) => {
            assign(lhs, rhs, bytecode, walk_context, &expr.span)?;
        }
        ExprKind::Variable(id) => {
            if walk_context.is_local {
                bytecode.byte_and_line(OpGetLocal as u8, expr.span.line);
            } else {
                bytecode.byte_and_line(OpGetGlobal as u8, expr.span.line);
            }
            bytecode.byte_and_line(*id, expr.span.line);
        }
        ExprKind::Tuple(exprs) => {
            for expr in exprs {
                generate_bytecode(walk_context, bytecode, expr)?;
            }
            bytecode.byte_and_line(OpMktuple as u8, expr.span.line);
            bytecode.byte_and_line(exprs.len() as u8, expr.span.line);
        }
        ExprKind::FunCall(name, params) => {
            for param in params {
                generate_bytecode(walk_context, bytecode, &param.value)?;
            }
            bytecode.byte_and_line(OpConst as u8, expr.span.line);
            bytecode.byte_and_line(*name, expr.span.line);
            bytecode.byte_and_line(OpCall as u8, expr.span.line);
            bytecode.byte_and_line(params.len() as u8, expr.span.line);
        }
        ExprKind::ValueDiscard => {
            bytecode.byte_and_line(OpPop as u8, expr.span.line);
        }
        ExprKind::FieldAccess(lhs, rhs) => {
            generate_bytecode(walk_context, bytecode, lhs)?;
            bytecode.byte(OpGetField as u8);
            bytecode.byte(*rhs);
        }
    }
    Ok(())
}

fn assign(
    lhs: &P<Expr>,
    rhs: &P<Expr>,
    bytecode: &mut Bytecode,
    walk_context: &mut WalkContext,
    span: &Span,
) -> Result<(), CompileError> {
    match &lhs.kind {
        ExprKind::Variable(id) => {
            generate_bytecode(walk_context, bytecode, rhs)?;
            if walk_context.is_local {
                bytecode.byte_and_line(OpSetLocal as u8, span.line);
            } else {
                bytecode.byte_and_line(OpSetGlobal as u8, span.line);
            }
            bytecode.byte_and_line(*id, span.line);
        }
        ExprKind::FunCall(name, parameters) => {
            for parameter in parameters {
                generate_bytecode(walk_context, bytecode, &parameter.value)?;
            }
            bytecode.byte_and_line(OpMktuple as u8, span.line);
            bytecode.byte_and_line(parameters.len() as u8, span.line);
            generate_bytecode(walk_context, bytecode, rhs)?;
            bytecode.byte_and_line(OpSetFunt as u8, span.line);
            bytecode.byte_and_line(*name, span.line);
        }
        ExprKind::Tuple(lhs_t) => match &rhs.kind {
            ExprKind::Tuple(rhs) => {
                if lhs_t.len() != rhs.len() {
                    return Err(CompileError::new(
                        format!("Tuple assignment requires target and source to have equal arities. Left was {}, right was {}", lhs_t.len(), rhs.len()),
                        lhs.span
                    ));
                }
                for i in 0..lhs_t.len() {
                    let left = lhs_t.get(i).unwrap();
                    let right = rhs.get(i).unwrap();
                    assign(left, right, bytecode, walk_context, &left.span)?;
                }
            }
            ExprKind::FunCall(_, _) => {
                // here we assume the function call will produce a tuple of the correct arity. This is not currently implemented
                // and will result in runtime error
                // it is also not possible to deconstruct nested tuples as results:
                // f(1) = (1, (2, 3))
                // (a, (b, c)) = f(1)
                // doesn't work with this code. This will be fixed later. Right now, nasty things will happen at runtime instead
                // Also, only variables are handled on lhs.
                generate_bytecode(walk_context, bytecode, rhs)?;
                for expr in lhs_t {
                    bytecode.byte_and_line(OpTupleGet as u8, rhs.span.line);
                    bytecode.byte_and_line(1, rhs.span.line);
                    match expr.kind {
                        ExprKind::Variable(id) => {
                            if walk_context.is_local {
                                bytecode.byte_and_line(OpSetLocal as u8, span.line);
                            } else {
                                bytecode.byte_and_line(OpSetGlobal as u8, span.line);
                            }
                            walk_context.initialized_variables.insert(id);
                            bytecode.byte_and_line(id, span.line);
                        }
                        _ => {
                            return Err(CompileError::new(
                                "deconstructing tuples only works with variables on lhs for now"
                                    .to_string(),
                                rhs.span,
                            ));
                        }
                    }
                    bytecode.byte(OpPop as u8);
                }
            }
            _ => {
                return Err(CompileError::new(
                    "Assigning to a tuple requires a tuple as the source".to_string(),
                    rhs.span,
                ));
            }
        },
        ExprKind::FieldAccess(path, name) => {
            generate_bytecode(walk_context, bytecode, path)?;
            generate_bytecode(walk_context, bytecode, rhs)?;
            bytecode.byte(OpSetField as u8);
            bytecode.byte(*name);
        }
        ExprKind::Literal(lit) => {
            return Err(CompileError::new(
                format!("Literal {} cannot be an assignment target", lit),
                lhs.span,
            ));
        }
        _ => {
            return Err(CompileError::new(
                format!("Expression {:?} cannot have its value assigned", lhs.kind),
                lhs.span,
            ));
        }
    }
    Ok(())
}

impl BinOpToken {
    fn bytecode(&self, bytecode: &mut Bytecode, line: usize) {
        match self {
            BinOpToken::Plus => {
                bytecode.byte_and_line(OpAdd as u8, line);
            }
            BinOpToken::Minus => {
                bytecode.byte_and_line(OpSubtract as u8, line);
            }
            BinOpToken::Star => {
                bytecode.byte_and_line(OpMultiply as u8, line);
            }
            BinOpToken::Slash => {
                bytecode.byte_and_line(OpDivide as u8, line);
            }
            BinOpToken::GreaterThan => {
                bytecode.byte_and_line(OpGt as u8, line);
            }
            BinOpToken::LessThan => {
                bytecode.byte_and_line(OpLt as u8, line);
            }
            BinOpToken::Equals => {
                bytecode.byte_and_line(OpEquals as u8, line);
            }
            BinOpToken::NotEquals => {
                bytecode.byte_and_line(OpEquals as u8, line);
                bytecode.byte_and_line(OpNot as u8, line);
            }
            BinOpToken::And => {
                bytecode.byte_and_line(OpAnd as u8, line);
            }
            BinOpToken::Or => {
                bytecode.byte_and_line(OpOr as u8, line);
            }
        }
    }
}

impl UnaryOpToken {
    fn bytecode(&self, bytecode: &mut Bytecode, line: usize) {
        match self {
            UnaryOpToken::Bang => {
                bytecode.byte_and_line(OpNot as u8, line);
            }
            UnaryOpToken::Minus => {
                bytecode.byte_and_line(OpNegate as u8, line);
            }
        }
    }
}

#[cfg(test)]
mod test {
    use crate::compiler::parser::ast_walker::walk;
    use crate::compiler::parser::tests::base_test;
    use crate::compiler::parser::CompileError;
    use crate::runtime::bytecode::OpCode::{OpAdd, OpConst, OpMultiply, OpSetGlobal};
    use crate::store::test_mock::StoreHasAllValuesMock;
    use crate::{integer_to_value, string_to_value};
    use std::collections::HashMap;

    impl From<CompileError> for Vec<CompileError> {
        fn from(value: CompileError) -> Self {
            vec![value]
        }
    }

    #[test]
    fn compile_binary_operator() -> Result<(), Vec<CompileError>> {
        let program = "1 + 2 * 3";

        let mut script = base_test(program, true)?;

        let mut store = StoreHasAllValuesMock {
            arities: HashMap::new(),
        };
        let bytecode = walk(&mut script, &mut store)?;

        assert_eq!(9, bytecode.script_bytecode.len());

        assert_eq!(OpConst as u8, bytecode.script_bytecode.get(0));
        assert_eq!(0_u8, bytecode.script_bytecode.get(1));
        assert_eq!(OpConst as u8, bytecode.script_bytecode.get(2));
        assert_eq!(1_u8, bytecode.script_bytecode.get(3));
        assert_eq!(OpConst as u8, bytecode.script_bytecode.get(4));
        assert_eq!(2_u8, bytecode.script_bytecode.get(5));
        assert_eq!(OpMultiply as u8, bytecode.script_bytecode.get(6));
        assert_eq!(OpAdd as u8, bytecode.script_bytecode.get(7));

        Ok(())
    }

    #[test]
    fn assign_integer_to_variable() -> Result<(), Vec<CompileError>> {
        let program = "a = 1";

        let mut script = base_test(program, true)?;
        let mut store = StoreHasAllValuesMock {
            arities: HashMap::new(),
        };
        let bytecode = walk(&mut script, &mut store)?;

        let mut values = vec![];
        values.push(integer_to_value!(1));

        assert_eq!(5, bytecode.script_bytecode.len());

        assert_eq!(OpConst as u8, bytecode.script_bytecode.get(0));
        assert_eq!(0_u8, bytecode.script_bytecode.get(1));
        assert_eq!(OpSetGlobal as u8, bytecode.script_bytecode.get(2));
        assert_eq!(0_u8, bytecode.script_bytecode.get(3));

        Ok(())
    }

    #[test]
    fn assign_value_to_tuple() -> Result<(), Vec<CompileError>> {
        let program = "f(1,3) = 2";

        let mut script = base_test(program, true)?;
        let mut store = StoreHasAllValuesMock {
            arities: HashMap::new(),
        };
        let bytecode = walk(&mut script, &mut store)?;

        let mut values = vec![];
        values.push(string_to_value!("f".to_string()));
        values.push(integer_to_value!(1));
        values.push(integer_to_value!(3));
        values.push(integer_to_value!(2));

        assert_eq!(11, bytecode.script_bytecode.len());

        Ok(())
    }
}