use crate::{codegen_binary, codegen_unary, opcodes::OpCode, parser::{ASTType, BinOp, Lit, UnOp, AST}, vm::{Value, VM}};

pub struct Codegen {
    pub vm: VM
}

impl Codegen {
    pub fn new(text: String) -> Self {
        Codegen {
            vm: VM::new(Vec::new(), Vec::new(),
                Vec::new(), text.lines().map(|s| s.to_string()).collect()),
        }
    }
    pub fn visit(&mut self, ast: AST) {
        use crate::parser::ASTType::*;
        match ast.ast_type {
            Program(asts) => {
                for ast in asts {
                    self.visit(ast);
                }
            }
            If(condition, block, else_expr) => {
                self.visit_if(*condition, *block, *else_expr);
            }
            Block(stmts) => {
                self.visit_block(stmts);
            }
            Expr(expr) => {
                self.visit_expr_stmt(*expr);
            }
            // REMOVE LATER
            Print(ref expr) => {
                self.visit_print(&ast, *expr.clone());
            }
            Binary(left, oper, right) => {
                self.visit_binary(*left, oper, *right);
            }
            Unary(oper, right) => {
                self.visit_unary(oper, *right);
            }
            Lit(ref lit) => {
                self.visit_lit(&ast, lit.clone());
            }
            _ => {}
        }
    }
    pub fn visit_if(&mut self, condition: AST, block: AST, else_expr: Option<AST>) {
        self.visit(condition);
        self.vm.chunk.push(OpCode::IF_NE.into());
        self.vm.chunk.push(0xFF); // placeholder
        self.vm.chunk.push(0xFF); // placeholder
        let start_jmp_loc = self.vm.chunk.len()-2;
        self.visit(block);
        if let Some(ast) = else_expr {
            match ast.ast_type {
                ASTType::If(_, _, _) => {
                    self.vm.chunk.push(OpCode::JMP.into());
                    self.vm.chunk.push(0xFF); // placeholder
                    self.vm.chunk.push(0xFF); // placeholder
                    let end_jmp_loc = self.vm.chunk.len()-2;
                    
                    self.patch_jump(start_jmp_loc);
                    self.visit(ast);
                    self.patch_jump(end_jmp_loc);
                }
                ASTType::Block(_) => {

                    self.vm.chunk.push(OpCode::JMP.into());
                    self.vm.chunk.push(0xFF); // placeholder
                    self.vm.chunk.push(0xFF); // placeholder
                    
                    let end_jmp_loc = self.vm.chunk.len()-2;
                    self.patch_jump(start_jmp_loc); // jmp to else
                    self.visit(ast);
                    self.patch_jump(end_jmp_loc); // jmp to after else
                }
                _ => {}
            }
        } else {
            self.patch_jump(start_jmp_loc);
        }
    }
    pub fn visit_block(&mut self, stmts: Vec<AST>) {
        for stmt in stmts {
            self.visit(stmt);
        }
    }
    pub fn visit_expr_stmt(&mut self, expr: AST) {
        self.visit(expr);
        self.vm.chunk.push(OpCode::POP.into());
    }
    // REMOVE LATER
    pub fn visit_print(&mut self, ast: &AST, expr: AST) {
        self.visit(expr);
        self.vm.chunk.push(OpCode::PRINT.into());
        if self.vm.lines.last().unwrap().0 != ast.end_line {
            self.vm.lines.push((ast.end_line, 1));
        } else {
            self.vm.lines.last_mut().unwrap().1 += 1;
        }
    }
    pub fn visit_binary(&mut self, left: AST, oper: BinOp, right: AST) {
        use crate::parser::BinOp::*;
        codegen_binary!(self, left, right, oper, Add, ADD, Sub, SUB, Mul, MUL,
            Div, DIV, Mod, MOD, BitAnd, BITAND, BitXor, BITXOR, BitOr, BITOR,
            Eq, EQ, Ne, NE, Gr, GR, Ge, GE, Ls, LS, Le, LE, Or, OR, And, AND);
    }
    pub fn visit_unary(&mut self, oper: UnOp, right: AST) {
        use crate::parser::UnOp::*;
        codegen_unary!(self, right, oper, Not, NOT, Neg, NEG);
    }
    pub fn visit_lit(&mut self, ast: &AST, lit: Lit) {
        match lit {
            Lit::Num(value) => {
                self.vm.pool.push(Value::Num(value));
                self.vm.chunk.push(OpCode::CONSTANT.into());
                self.vm.chunk.push((self.vm.pool.len()-1) as u8);
                if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
                    self.vm.lines.push((ast.end_line, 2));
                } else {
                    self.vm.lines.last_mut().unwrap().1 += 2;
                }
            }
            Lit::Str(value) => {
                self.vm.pool.push(Value::Str(value.into()));
                self.vm.chunk.push(OpCode::CONSTANT.into());
                self.vm.chunk.push((self.vm.pool.len()-1) as u8);
                if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
                    self.vm.lines.push((ast.end_line, 2));
                } else {
                    self.vm.lines.last_mut().unwrap().1 += 2;
                }
            }
            Lit::Bool(value) => {
                self.vm.pool.push(Value::Bool(value));
                self.vm.chunk.push(OpCode::CONSTANT.into());
                self.vm.chunk.push((self.vm.pool.len()-1) as u8);
                if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
                    self.vm.lines.push((ast.end_line, 2));
                } else {
                    self.vm.lines.last_mut().unwrap().1 += 2;
                }
            }
            _ => {}
        }
    }
    pub fn patch_jump(&mut self, loc: usize) {
        let jump = self.vm.chunk.len() - loc - 2;

        self.vm.chunk[loc] = ((jump >> 8) & 0xFF) as u8;
        self.vm.chunk[loc + 1] = (jump & 0xFF) as u8;
    }
}