import io:*
import mem:*
import map:*
import hash_map:*
import stack:*
import string:*
import util:*
import tree:*
import symbol:*
import ast_nodes:*
// for error with syntax tree
import pass_common:*
import poset:*


adt byte_inst {
    nop,
    imm: imm,
    alloca: alloca,
    ld: ld,
    st: st,
    call,
    ret: ret
}
obj imm {
    var reg: int
    var val: int
}
obj alloca {
    var reg: int
    var ident: *ast_node
}
obj ld {
    var to_reg: int
    var from_reg: int
    var ident: *ast_node
}
obj st {
    var to_reg: int
    var from_reg: int
}
obj ret {
    var reg: int
}

fun to_string(b: byte_inst): string {
    match (b) {
        byte_inst::nop()     return string("nop")
        byte_inst::imm(i)    return string("r") + i.reg + " = imm " + i.val
        byte_inst::alloca(a) return string("r") + a.reg + " = alloca(" + a.ident->identifier.name + ")"
        byte_inst::ld(l)     return string("r") + l.to_reg + " = ld r" + l.from_reg + " (" + l.ident->identifier.name + ")"
        byte_inst::st(s)     return string("st(r") + s.to_reg + " <= r" + s.from_reg + ")"
        byte_inst::call()    return string("call")
        byte_inst::ret(r)    return string("ret r") + r.reg
    }
    return string("Missed byte_inst case in to_string")
}

fun bytecode_to_string(bytecode: ref vector<basic_block>): string {
    return string("\n").join(bytecode.map(fun(bb: ref basic_block): string return bb.to_string();))
}

fun basic_block(name: ref string): basic_block {
    var to_ret.construct(name): basic_block
    return to_ret
}
obj basic_block (Object) {
    var name: string
    var instructions: vector<byte_inst>

    fun construct(): *basic_block {
        instructions.construct()
        name.construct()
        return this
    }
    fun construct(name_in: ref string): *basic_block {
        instructions.construct()
        name.copy_construct(&name_in)
        return this
    }
    fun copy_construct(old: *basic_block) {
        instructions.copy_construct(&old->instructions)
        name.copy_construct(&old->name)
    }
    fun operator=(other: ref basic_block) {
        destruct()
        copy_construct(&other)
    }
    fun destruct() {
        instructions.destruct()
        name.destruct()
    }
    fun to_string(): string {
        var res = name + ":\n"
        instructions.for_each(fun(b: byte_inst) {
            res += "\t" + to_string(b) + "\n"
        })
        return res
    }
}

obj bytecode_generator (Object) {
    var reg_counter: int
    var id_counter: int
    var ast_name_map: hash_map<*ast_node, string>
    var var_to_reg: map<*ast_node, int>
    var blocks: vector<basic_block>
    fun construct(): *bytecode_generator {
        reg_counter = 0
        id_counter = 0
        ast_name_map.construct()
        var_to_reg.construct()
        blocks.construct()

        return this
    }
    fun copy_construct(old: *bytecode_generator) {
        reg_counter = old->reg_counter
        id_counter = old->id_counter
        ast_name_map.copy_construct(&old->ast_name_map)
        var_to_reg.copy_construct(&old->var_to_reg)
        blocks.copy_construct(&old->blocks)
    }
    fun operator=(other: ref bytecode_generator) {
        destruct()
        copy_construct(&other)
    }
    fun destruct() {
        ast_name_map.destruct()
        var_to_reg.destruct()
        blocks.destruct()
    }
    fun get_id(): string return to_string(id_counter++);
    fun get_reg(): int return reg_counter++;
    fun generate_bytecode(name_ast_map: map<string, pair<*tree<symbol>,*ast_node>>): vector<basic_block> {

        // iterate through asts
        name_ast_map.for_each(fun(name: string, tree_pair: pair<*tree<symbol>,*ast_node>) {
            // iterate through children for each ast
            // do lambdas seperatly, so we can reconstitute the enclosing object if it has one
            tree_pair.second->translation_unit.lambdas.for_each(fun(child: *ast_node) {
                generate_function_definition(child)
            })
            tree_pair.second->translation_unit.children.for_each(fun(child: *ast_node) {
                match (*child) {
                    ast_node::declaration_statement(backing)    generate_declaration_statement(child)
                    ast_node::compiler_intrinsic(backing)       generate_compiler_intrinsic(child)
                    ast_node::function(backing)                 generate_function_definition(child)
                    ast_node::template(backing) {
                        backing.instantiated.for_each(fun(node: *ast_node) {
                            match (*node) {
                                ast_node::function(backing) generate_function_definition(node)
                                ast_node::type_def(backing) {
                                    backing.methods.for_each(fun(method: *ast_node) {
                                        if (is_template(method))
                                            method->template.instantiated.for_each(fun(m: *ast_node) generate_function_definition(m);)
                                        else
                                            generate_function_definition(method)
                                    })
                                }
                            }
                        })
                    }
                    ast_node::type_def(backing) {
                        backing.methods.for_each(fun(method: *ast_node) {
                            if (is_template(method))
                                method->template.instantiated.for_each(fun(m: *ast_node) generate_function_definition(m);)
                            else
                                generate_function_definition(method)
                        })
                    }
                }
            })
        })
        return blocks
    }
    fun generate_function_definition(node: *ast_node): int {
        blocks.add(basic_block(get_name(node)))
        node->function.parameters.for_each(fun(p: *ast_node) {
            var_to_reg[p] = emit_alloca(p)
        })
        generate(node->function.body_statement)
        return -1
    }
    fun generate_declaration_statement(node: *ast_node): int {
        var identifier = node->declaration_statement.identifier
        var ident_type = identifier->identifier.type
        var_to_reg[identifier] = emit_alloca(identifier)
        if (node->declaration_statement.expression) {
            emit_st(var_to_reg[identifier], generate(node->declaration_statement.expression))
        }
        return -1
    }
    fun generate_assignment_statement(node: *ast_node): int {
        var to   = generate(node->assignment_statement.to, true)
        var from = generate(node->assignment_statement.from)
        emit_st(to, from)
        return -1
    }
    fun generate_if_statement(node: *ast_node): int {
        generate(node->if_statement.condition)
        generate(node->if_statement.then_part)
        if (node->if_statement.else_part)
            generate(node->if_statement.else_part)
        return -1
    }
    fun generate_while_loop(node: *ast_node): int {
        generate(node->while_loop.condition)
        generate(node->while_loop.statement)
        return -1
    }
    fun generate_for_loop(node: *ast_node): int {
        if (node->for_loop.init)
            generate(node->for_loop.init)
        if (node->for_loop.condition)
            generate(node->for_loop.condition)
        if (node->for_loop.update)
            generate(node->for_loop.update)
        generate(node->for_loop.body)
        return -1
    }
    fun generate_identifier(node: *ast_node, lvalue: bool): int {
        if (lvalue)
            return var_to_reg[node]
        else
            return emit_ld(node)
    }
    fun generate_return_statement(node: *ast_node): int {
        if (node->return_statement.return_value)
            emit_ret(generate(node->return_statement.return_value))
        else
            emit_ret(-1)
        return -1
    }
    fun generate_branching_statement(node: *ast_node): int {
        match(node->branching_statement.b_type) {
            branching_type::break_stmt()    blocks.last().instructions.add(byte_inst::nop())
            branching_type::continue_stmt() blocks.last().instructions.add(byte_inst::nop())
        }
        return -1
    }
    fun generate_cast(node: *ast_node): int {
        return generate(node->cast.value)
    }
    fun generate_value(node: *ast_node): int {
        return emit_imm(string_to_num<int>(node->value.string_value))
    }
    fun generate_code_block(node: *ast_node): int {
        node->code_block.children.for_each(fun(child: *ast_node) generate(child);)
        return -1
    }
    // this generates the function as a value, not the actual function
    fun generate_function(node: *ast_node): int {
        return emit_imm(-2)
    }
    fun generate_function_call(node: *ast_node, lvalue: bool): int {
        node->function_call.parameters.for_each(fun(child: *ast_node) generate(child);)
        return emit_call()
    }

    fun generate_compiler_intrinsic(node: *ast_node): int {
        blocks.last().instructions.add(byte_inst::nop())
        return -1
    }

    fun generate(node: *ast_node): int return generate(node, false)
    fun generate(node: *ast_node, lvalue: bool): int {
        match (*node) {
            ast_node::declaration_statement(backing)    return generate_declaration_statement(node)
            ast_node::assignment_statement(backing)     return generate_assignment_statement(node)
            ast_node::if_statement(backing)             return generate_if_statement(node)
            ast_node::while_loop(backing)               return generate_while_loop(node)
            ast_node::for_loop(backing)                 return generate_for_loop(node)
            ast_node::function(backing)                 return generate_function(node)
            ast_node::function_call(backing)            return generate_function_call(node, lvalue)
            ast_node::compiler_intrinsic(backing)       return generate_compiler_intrinsic(node)
            ast_node::code_block(backing)               return generate_code_block(node)
            ast_node::return_statement(backing)         return generate_return_statement(node)
            ast_node::branching_statement(backing)      return generate_branching_statement(node)
            ast_node::cast(backing)                     return generate_cast(node)
            ast_node::value(backing)                    return generate_value(node)
            ast_node::identifier(backing)               return generate_identifier(node, lvalue)
        }
        error("Bad node")
    }
    fun get_name(node: *ast_node): string {
        var maybe_it = ast_name_map.get_ptr_or_null(node);
        if (maybe_it)
            return *maybe_it
        var result = get_ast_name(node) + get_id()
        if (is_function(node) && node->function.name == "main")
            result = "main"
        ast_name_map.set(node, result)
        return result
    }
    fun emit_imm(value: int): int {
        var i: imm
        i.reg = get_reg()
        i.val = value
        blocks.last().instructions.add(byte_inst::imm(i))
        return i.reg
    }
    fun emit_alloca(node: *ast_node): int {
        var a: alloca
        a.reg = get_reg()
        a.ident = node
        blocks.last().instructions.add(byte_inst::alloca(a))
        return a.reg
    }
    fun emit_ld(node: *ast_node): int {
        var l: ld
        l.to_reg = get_reg()
        l.from_reg = var_to_reg[node]
        l.ident = node
        blocks.last().instructions.add(byte_inst::ld(l))
        return l.to_reg
    }
    fun emit_st(to_reg: int, from_reg: int): int {
        var s: st
        s.to_reg = to_reg
        s.from_reg = from_reg
        blocks.last().instructions.add(byte_inst::st(s))
        return -1
    }
    fun emit_ret(reg: int): int {
        var r: ret
        r.reg = reg
        blocks.last().instructions.add(byte_inst::ret(r))
        return -1
    }
    fun emit_call(): int {
        blocks.last().instructions.add(byte_inst::call())
        return -1
    }

    fun evaluate(): int {
        println("evaling main")
        var main_entry = blocks.find_first_satisfying(fun(block: basic_block): bool return block.name == "main";)
        var registers.construct(reg_counter): vector<int>
        registers.size = reg_counter
        var stack_mem.construct(): vector<int>
        for (var i = 0; i < main_entry.instructions.size; i++;) {
            match(main_entry.instructions[i]) {
                byte_inst::nop()     {}
                byte_inst::imm(i)    registers[i.reg] = i.val
                byte_inst::alloca(a) { stack_mem.add(0); registers[a.reg] = stack_mem.size -1; }
                byte_inst::ld(l)     registers[l.to_reg] = stack_mem[registers[l.from_reg]]
                byte_inst::st(s)     stack_mem[registers[s.to_reg]] = registers[s.from_reg]
                byte_inst::call()    error("call")
                byte_inst::ret(r)    return registers[r.reg]
            }
        }
        return -1
    }
}