use crate::prelude::*;

use Instruction as I;
use Value as V;
use ExprData as E;

pub type NamesTable = Rc<RefCell<Vec<Rc<str>>>>;
pub type GlobalsTable = Rc<RefCell<Vec<Global>>>;

pub struct CompilerBuilder<'c> {
    globals: GlobalsTable,
    names: NamesTable,
    repl: bool,
    debug: bool,
    name: Rc<str>,
    parent: Option<&'c Compiler<'c>>,
}

impl<'c> CompilerBuilder<'c> {

    pub fn new() -> Self {
        Self {
            globals: Rc::new(RefCell::new(Vec::new())),
            repl: false,
            debug: false,
            name: "<root>".into(),
            parent: None,
            names: Rc::new(RefCell::new(Vec::new()))
        }
    }

    pub fn repl(mut self, repl: bool) -> Self {
        self.repl = repl;
        self
    }

    pub fn debug(mut self, debug: bool) -> Self {
        self.debug = debug;
        self
    }

    pub fn globals(mut self, globals: GlobalsTable) -> Self {
        self.globals = globals;
        self
    }

    pub fn names(mut self, names: NamesTable) -> Self {
        self.names = names;
        self
    }

    pub fn parent(mut self, parent: &'c Compiler) -> Self {
        self.parent = Some(parent);
        self
    }

    pub fn name(mut self, name: Rc<str>) -> Self {
        self.name = name;
        self
    }

    pub fn build(self) -> Compiler<'c> {
        Compiler {
            name: self.name,
            names: self.names,
            parent: self.parent,
            globals: self.globals,
            repl: self.repl,
            debug: self.debug,
            scopes: Vec::new(),
            locals: Vec::new(),
            chunk: Chunk::new(),
            loop_top: Vec::new(),
            loop_bot: Vec::new(),
            root_is_block: false,
        }
    }
}

pub struct Compiler<'c> {
    name: Rc<str>,
    parent: Option<&'c Compiler<'c>>,

    locals: Vec<Local>,
    globals: GlobalsTable,
    names: NamesTable,

    root_is_block: bool,

    scopes: Vec<usize>,
    chunk: Chunk,

    loop_top: Vec<(usize, usize)>,
    loop_bot: Vec<usize>,

    repl: bool,
    debug: bool,
}

#[derive(Clone)]
struct Local {
    name: Rc<str>,
    idx: usize,
    scope: usize,
    is_const: bool,
}

#[derive(Clone)]
pub struct Global {
    pub name: Rc<str>,
    pub idx: usize,
    pub is_const: bool,
}

macro_rules! error {
    ($($arg:tt)*) => {
        exception!(COMPILE_EXCEPTION, $($arg)*)
    };
}

impl<'c> Compiler<'c> {

    fn child(&'c self, name: Rc<str>) -> Self {
        CompilerBuilder::new()
            .name(name)
            .debug(self.debug)
            .globals(self.globals.clone())
            .names(self.names.clone())
            .repl(false)
            .parent(self)
            .build()
    }

    fn begin_scope(&mut self) {
        if self.root_is_block {
            self.root_is_block = false;
        } else {
            self.scopes.push(self.locals.len())
        }
    }

    fn collapse_scopes(&mut self, top_scope: usize) {
        let mut cutoff = usize::MAX;
        while self.scopes.len() > top_scope {
            cutoff = self.scopes.pop().expect("bypassed compiler scope check")
        }
        if cutoff < self.locals.len() {
            self.emit(Instruction::DiscardLocals((self.locals.len() - cutoff) as u16), self.last_pos());
            self.locals.truncate(cutoff);
        };
    }

    fn end_scope(&mut self) {
        let Some(cutoff) = self.scopes.pop() else {
            return;
        };
        if cutoff < self.locals.len() {
            self.emit(Instruction::DiscardLocals((self.locals.len() - cutoff) as u16), self.last_pos());
            self.locals.truncate(cutoff);
        };
    }

    fn create_local(&mut self, name: Rc<str>, is_const: bool) -> Local {
        let local = Local { name, idx: self.locals.len(), scope: self.scopes.len(), is_const };
        self.locals.push(local.clone());
        local
    }

    fn create_global(&mut self, name: Rc<str>, is_const: bool) -> Global {
        let global = Global { name, idx: self.globals.borrow().len(), is_const };
        self.globals.borrow_mut().push(global.clone());
        global
    }

    fn create_local_checked(&mut self, name: Rc<str>, is_const: bool, pos: Position) -> Result<Local> {
        if let Some(local) = self.find_local(&name) {
            if local.scope == self.scopes.len() {
                return Err(error!("redefined {name}").pos(pos))
            }
        };
        Ok(self.create_local(name, is_const))
    }

    fn create_global_checked(&mut self, name: Rc<str>, is_const: bool, pos: Position) -> Result<Global> {
        if let Some(_) = self.find_global(&name) {
            return Err(error!("redefined {name}").pos(pos))
        }
        Ok(self.create_global(name, is_const))
    }

    fn find_local(&self, name: &str) -> Option<Local> {
        for local in self.locals.iter().rev() {
            if local.name.as_ref() == name {
                return Some(local.clone())
            }
        }
        None
    }

    fn find_global(&self, name: &str) -> Option<Global> {
        if let Some(parent) = self.parent {
            return parent.find_global(name)
        }
        for global in self.globals.borrow().iter() {
            if global.name.as_ref() == name {
                return Some(global.clone())
            }
        }
        None
    }

    fn get_name(&mut self, name: Rc<str>) -> usize {
        // TODO: find name if already exists
        let idx = self.names.borrow().len();
        self.names.borrow_mut().push(name);
        idx
    }

    fn emit_const(&mut self, val: Value, pos: Position) {
        // TODO: find constant if already exists
        self.chunk.constants.push(val);
        let c = self.chunk.constants.len() - 1;
        self.emit(Instruction::Const(c as u16), pos);
    }

    fn can_make_globals(&self) -> bool {
        self.repl && self.parent.is_none() && self.scopes.len() == 0
    }

    fn compile_value(&mut self, val: &Value, pos: Position) {
        match val {
            V::Nil => self.emit(I::Nil, pos),
            V::Bool(b) => if *b { self.emit(I::True, pos) } else { self.emit(I::False, pos) },
            V::Int(i) => {
                if let Ok(i) = i16::try_from(*i) {
                    self.emit(I::Int(i), pos);
                } else {
                    self.emit_const(val.clone(), pos);
                }
            },
            _ => self.emit_const(val.clone(), pos),
        }
    }



    fn finish_loop(&mut self) {
        self.loop_top.pop();
        while let Some(tmp) = self.loop_bot.pop() {
            self.re_emit(I::Jump(self.cur()), tmp as usize);
        }
    }

    fn compile_expr(&mut self, expr: &Expr) -> Result<()> {
        match &expr.data {
            E::NoOp => self.emit(I::Nil, expr.pos),
            E::If(cond, ifb, elseb) => {
                self.compile_expr(cond)?;
                let jmpidx = self.emit_temp(expr.pos);
                self.compile_expr(ifb)?;
                let jmpidx2 = self.emit_temp(expr.pos);
                self.re_emit(I::JumpFalse(self.cur()), jmpidx);
                if let Some(elseb) = elseb {
                    self.compile_expr(elseb)?;
                }
                self.re_emit(I::Jump(self.cur()), jmpidx2);
            },
            E::Function(name, params, body, varadic) => {
                let chunk = self.compile_function(name.clone(), params, body)?;
                let arity = params.len() - if *varadic { 1 } else { 0 };
                let fun = Value::Function(Rc::new(
                        Function {
                            name: name.0.clone(),
                            arity,
                            fun: InnerFunction::Compiled(chunk.into()),
                            variadic: *varadic
                        }
                ));
                self.emit_const(fun, expr.pos);
                self.emit(I::Dup, expr.pos);
                if self.can_make_globals() {
                    let global = self.create_global_checked(name.0.clone(), false, name.1)?;
                    self.emit(I::StoreGlobal(global.idx as u16), expr.pos);
                } else {
                    self.create_local_checked(name.0.clone(), false, name.1)?;
                    self.emit(I::CreateLocal, expr.pos);
                }
            },
            E::Lambda(params, body, varadic) => {
                let name: AstName = (Rc::from("<lambda>"), expr.pos);
                let chunk = self.compile_function(name.clone(), params, body)?;
                let arity = params.len() - if *varadic { 1 } else { 0 };
                let fun = Value::Function(Rc::new(
                        Function {
                            name: name.0.clone(),
                            arity,
                            fun: InnerFunction::Compiled(chunk.into()),
                            variadic: *varadic
                        }
                ));
                self.emit_const(fun, expr.pos);
            },
            E::Loop(expr) => {
                let idx = self.cur();
                self.loop_top.push((idx as usize, self.scopes.len()));
                self.compile_expr(expr)?;
                self.emit(I::Discard(1), expr.pos);
                self.emit(I::Jump(idx), expr.pos);
                self.finish_loop();
                self.emit(I::Nil, expr.pos);
            },
            E::Try(expr, err, catch) => {
                let jmpidx = self.emit_temp(expr.pos);
                self.compile_expr(expr)?;
                self.emit(I::TryEnd, expr.pos);
                let jmpidx2 = self.emit_temp(expr.pos);
                self.re_emit(I::Try(self.cur()), jmpidx);
                self.begin_scope();
                self.create_local(err.0.clone(), true);
                self.emit(I::CreateLocal, err.1);
                self.compile_expr(catch)?;
                self.end_scope();
                self.re_emit(I::Jump(self.cur() as u16), jmpidx2);
            },
            E::While(cond, expr) => {
                let top = self.cur();
                self.compile_expr(cond)?;
                let jmpidx = self.emit_temp(expr.pos);
                self.loop_top.push((top as usize, self.scopes.len()));
                self.compile_expr(expr)?;
                self.emit(I::Discard(1), expr.pos);
                self.emit(I::Jump(top), expr.pos);
                self.re_emit(I::JumpFalse(self.cur()), jmpidx);
                self.finish_loop();
                self.emit(I::Nil, expr.pos);
            },
            E::DoWhile(expr, cond) => {
                let top = self.cur();
                self.loop_top.push((top as usize, self.scopes.len()));
                self.compile_expr(expr)?;
                self.emit(I::Discard(1), expr.pos);
                self.compile_expr(cond)?;
                self.emit(I::JumpTrue(top), expr.pos);
                self.finish_loop();
                self.emit(I::Nil, expr.pos);
            },
            E::For(name, cond, expr) => {
                self.compile_expr(cond)?;
                self.emit(I::IterCreate, cond.pos);
                let top = self.cur();
                self.emit(I::IterNext, cond.pos);
                self.emit(I::Dup, expr.pos);
                let jumpidx = self.emit_temp(expr.pos);
                self.loop_top.push((top as usize, self.scopes.len()));
                self.begin_scope();
                self.create_local(name.0.clone(), true);
                self.emit(I::CreateLocal, name.1);
                self.compile_expr(expr)?;
                self.emit(I::Discard(1), expr.pos);
                self.end_scope();
                self.emit(I::Jump(top), expr.pos);
                self.re_emit(I::JumpNil(self.cur()), jumpidx);
                self.finish_loop();
                self.emit(I::Discard(2), expr.pos);
                self.emit(I::Nil, expr.pos);
            }
            E::Block(block) => {
                if block.len() == 0 {
                    self.emit(I::Nil, expr.pos);
                    return Ok(());
                }
                self.begin_scope();
                for (i, expr) in block.iter().enumerate() {
                    self.compile_expr(expr)?;
                    if i + 1 != block.len() {
                        self.emit(I::Discard(1), expr.pos);
                    }
                }
                self.end_scope();
            },
            E::Let(name, expr) => {
                self.compile_expr(expr)?;
                self.emit(I::Dup, expr.pos);
                if self.can_make_globals() {
                    let global = self.create_global_checked(name.0.clone(), false, name.1)?;
                    self.emit(I::StoreGlobal(global.idx as u16), expr.pos);
                } else {
                    self.create_local_checked(name.0.clone(), false, name.1)?;
                    self.emit(I::CreateLocal, expr.pos);
                }
            },
            E::Const(name, expr) => {
                self.compile_expr(expr)?;
                self.emit(I::Dup, expr.pos);
                if self.can_make_globals() {
                    let global = self.create_global_checked(name.0.clone(), true, name.1)?;
                    self.emit(I::StoreGlobal(global.idx as u16), expr.pos);
                } else {
                    self.create_local_checked(name.0.clone(), true, name.1)?;
                    self.emit(I::CreateLocal, expr.pos);
                }
            },
            E::Continue => {
                let top = self.loop_top.pop();
                if let Some((top, scope)) = top {
                    self.collapse_scopes(scope);
                    self.emit(I::Jump(top as u16), expr.pos);
                } else {
                    return Err(error!("invalid continue outside loop").pos(expr.pos))
                }
            },
            E::Break => {
                let top = self.loop_top.pop();
                if let Some((_, scope)) = top {
                    self.collapse_scopes(scope);
                    let tmpidx = self.emit_temp(expr.pos);
                    self.loop_bot.push(tmpidx);
                } else {
                    return Err(error!("invalid break outside loop").pos(expr.pos))
                }
            },
            E::Return(expr) => {
                self.compile_expr(expr)?;
                self.emit(I::Return, expr.pos);
            },
            E::Literal(val) => self.compile_value(val, expr.pos),
            E::Ident(name) => {
                if name.as_ref() == "_" {
                    self.emit(I::Nil, expr.pos);
                } else if let Some(local) = self.find_local(name) {
                    self.emit(I::LoadLocal(local.idx as u16), expr.pos);
                } else if let Some(global) = self.find_global(name) {
                    self.emit(I::LoadGlobal(global.idx as u16), expr.pos);
                } else {
                    return Err(error!("variable '{name}' is undefined").pos(expr.pos))
                };
            },
            E::Assign(lhs, rhs) => {
                self.compile_expr(rhs)?;
                self.emit(I::Dup, rhs.pos);
                match &lhs.data {
                    E::Ident(name) if name.as_ref() != "_" => {
                        if let Some(local) = self.find_local(&name) {
                            if local.is_const {
                                return Err(error!("cannot assign to const '{name}'").pos(lhs.pos))
                            }
                            self.emit(I::StoreLocal(local.idx as u16), lhs.pos);
                        } else if let Some(global) = self.find_global(&name) {
                            if global.is_const {
                                return Err(error!("cannot assign to const '{name}'").pos(lhs.pos))
                            }
                            self.emit(I::StoreGlobal(global.idx as u16), lhs.pos);
                        } else if self.can_make_globals() {
                            let global = self.create_global_checked(name.clone(), false, lhs.pos)?;
                            self.emit(I::StoreGlobal(global.idx as u16), lhs.pos);
                        } else {
                            self.create_local_checked(name.clone(), false, lhs.pos)?;
                            self.emit(I::CreateLocal, lhs.pos);
                        }
                    },
                    E::Index(expr, params) => {
                        self.compile_expr(expr)?;
                        for param in params {
                            self.compile_expr(param)?;
                        }
                        self.emit(I::StoreIndex(params.len() as u8), expr.pos);
                    },
                    E::FieldAccess(expr, ident) => {
                        self.compile_expr(expr)?;
                        let name = self.get_name(ident.0.clone());
                        self.emit(I::StoreField(name as u16), expr.pos);
                    }
                    _ => return Err(error!("assignment to {lhs} is not allowed").pos(lhs.pos))
                }
            }
            E::UnaryOp(expr, op) => {
                self.compile_expr(expr)?;
                self.emit(I::UnaryOp(*op), expr.pos);
            },
            E::BinaryOp(lhs, rhs, op) => {
                self.compile_expr(lhs)?;
                self.compile_expr(rhs)?;
                self.emit(I::BinaryOp(*op), lhs.pos);
            },
            E::Index(expr, params) => {
                self.compile_expr(expr)?;
                for param in params {
                    self.compile_expr(param)?;
                }
                self.emit(I::Index(params.len() as u8), expr.pos)
            },
            E::FnCall(fun, params) => {
                for expr in params {
                    self.compile_expr(expr)?;
                }
                self.compile_expr(fun)?;
                self.emit(I::Call(params.len() as u8), expr.pos);
            },
            E::FieldAccess(expr, field) => {
                self.compile_expr(expr)?;
                let idx = self.get_name(field.0.clone());
                self.emit(I::Field(idx as u16), expr.pos)
            }
            E::List(list) => {
                for expr in list {
                    self.compile_expr(expr)?;
                }
                self.emit(I::NewList(list.len() as u16), expr.pos);
            },
            E::Matrix(mat) => {
                for expr in &mat.2 {
                    self.compile_expr(expr)?;
                }
                self.emit(I::NewMatrix(mat.2.len() as u16, mat.0 as u8), expr.pos);
            },
            E::Table(table) => {
                for (key, value) in table {
                    self.compile_expr(key)?;
                    self.compile_expr(value)?;
                }
                self.emit(I::NewTable(table.len() as u16), expr.pos);
            },
            E::And(lhs, rhs) => {
                self.compile_expr(lhs)?;
                self.emit(I::Dup, lhs.pos);
                let jmpidx = self.emit_temp(lhs.pos);
                self.compile_expr(rhs)?;
                self.re_emit(I::JumpFalse(self.cur()), jmpidx);
            },
            E::Or(lhs, rhs) => {
                self.compile_expr(lhs)?;
                self.emit(I::Dup, lhs.pos);
                let jmpidx = self.emit_temp(lhs.pos);
                self.compile_expr(rhs)?;
                self.re_emit(I::JumpTrue(self.cur()), jmpidx);
            },
            E::Pipeline(lhs, rhs) => {
                self.compile_expr(lhs)?;
                self.compile_expr(rhs)?;
                self.emit(I::Call(1), expr.pos);
            }
        };
        Ok(())
    }

    fn compile_function(&mut self, name: AstName, params: &Vec<AstName>, body: &Box<Expr>) -> Result<Chunk> {
        let mut compiler = self.child(name.0);
        for (name, pos) in params {
            compiler.create_local(name.clone(), false);
            compiler.emit(I::CreateLocal, *pos);
        }
        compiler.compile_expr(body)?;
        compiler.finish(name.1)?;
        Ok(compiler.chunk)
    }

    fn cur(&self) -> u16 {
        self.chunk.code.len() as u16
    }

    fn emit_temp(&mut self, pos: Position) -> usize {
        let idx = self.chunk.code.len();
        self.emit(Instruction::NoOp, pos);
        idx
    }

    fn emit(&mut self, ins: Instruction, pos: Position) {
        //println!("{}: {ins}", self.name);
        self.chunk.code.push(ins);
        self.chunk.pos.push(pos);
    }

    fn re_emit(&mut self, ins: Instruction, idx: usize) {
        //println!("{} at {}: {ins}", self.name, &self.chunk.code[idx]);
        self.chunk.code[idx] = ins;
    }

    fn last_pos(&self) -> Position {
        if let Some(pos) = self.chunk.pos.last() {
            *pos
        } else {
            Position::default()
        }
    }

    fn finish(&mut self, pos: Position) -> Result<()> {
        let ins = match self.chunk.code.last() {
            Some(ins) => ins.clone(),
            None => {
                self.emit(I::Nil, pos);
                self.emit(I::Return, pos);
                if self.debug {
                    println!("{}\n{}", self.name, self.chunk);
                }
                return Ok(())
            }
        };
        match ins {
            I::Return => {},
            _ => {
                self.emit(I::Return, pos);
            }
        };

        if self.loop_bot.len() > 0 {
            return Err(error!("invalid break outside loop").pos(pos))
        }

        if self.debug {
            println!("{}\n{}", self.name, self.chunk);
        }
        Ok(())
    }

    pub fn compile(
        &mut self,
        body: &Expr,
    ) -> Result<Rc<Function>> {
        if let ExprData::Block(_) = &body.data {
            self.root_is_block = true;
        }
        self.chunk = Chunk::new();
        self.compile_expr(body)?;
        self.finish(self.last_pos())?;
        let fun = Function { name: self.name.clone(), fun: InnerFunction::Compiled(self.chunk.clone().into()), arity: 0, variadic: false };
        Ok(Rc::new(fun))
    }
}