pub mod basic_block;
pub mod error;
mod locals;
pub mod thunk;

use crate::{syn::ast::Body, obj::prelude::*, vm::{consts::*, inst::Inst}};
use std::collections::{BTreeMap, HashMap};

#[derive(Default)]
pub struct Compile {
    const_data: ConstData,
    globals: BTreeMap<Sym, Name>,
    scope: Vec<BTreeMap<Sym, Name>>,
    names: Vec<Sym>,
    next_local: Name,
}

impl Compile {
    /// Creates a new compiler using the given text.
    pub fn new() -> Self {
        Default::default()
    }

    /// Compiles the given AST body.
    pub fn compile<'c>(&'c mut self, body: &Body) -> error::Result<(Vec<Inst>, &'c ConstPool)> {
        let main = thunk::CompileBody::new(self)
            .compile(body)?
            .flatten()
            .to_vec();
        Ok((main, &self.const_data.const_pool))
    }

    /// Gets the constant data that is interned in this compile session.
    pub fn const_data(&self) -> &ConstData {
        &self.const_data
    }

    /// Mutably gets the constant data that is interned in this compile session.
    pub fn const_data_mut(&mut self) -> &mut ConstData {
        &mut self.const_data
    }

    /// Gets or inserts a static int reference.
    pub fn const_int(&mut self, int: i64) -> (ConstHandle, IntRef) {
        self.const_data_mut().const_int(int)
    }

    /// Gets or inserts a static string reference.
    pub fn const_str(&mut self, s: impl AsRef<str>) -> (ConstHandle, StrRef) {
        self.const_data_mut().const_str(s)
    }

    /// Looks up a variable name.
    ///
    /// This will search up the locals stack until the given name is found, ultimately ending with
    /// a global name lookup.
    pub fn lookup_scope(&mut self, sym: Sym) -> Option<Name> {
        self.scope
            .iter()
            .rev()
            .filter_map(|locals| locals.get(&sym))
            .next()
            .or_else(|| self.globals.get(&sym))
            .copied()
    }

    /// Looks up a variable name, or creates a global name if it doesn't exist.
    pub fn lookup_scope_or_create_global(&mut self, sym: Sym) -> Name {
        if let Some(local) = self.lookup_scope(sym) {
            local
        } else {
            self.create_global(sym)
        }
    }

    /// Creates a new local variable if it does not exist in the current local scope.
    pub(crate) fn create_local(&mut self, sym: Sym) -> Name {
        let locals = self.scope.last_mut().expect("scope");
        if let Some(local) = locals.get(&sym) {
            *local
        } else {
            // wish I could use mem::replace here, oh well
            let local = self.next_local;
            self.next_local = self.next_local.next();
            locals.insert(sym, local);
            self.names.push(sym);
            assert_eq!(self.names.len(), self.next_local.index());
            local
        }
    }

    /// Creates a new global variable if it does not exist in the current global scope.
    pub fn create_global(&mut self, sym: Sym) -> Name {
        if let Some(global) = self.globals.get(&sym) {
            *global
        } else {
            let global = self.next_local;
            self.next_local = self.next_local.next();
            self.globals.insert(sym, global);
            self.names.push(sym);
            assert_eq!(self.names.len(), self.next_local.index());
            global
        }
    }

    /// Pushes an empty scope layer of local variables.
    pub(crate) fn push_scope_layer(&mut self) {
        self.scope.push(Default::default());
    }

    /// Pops a scope layer of local variables, if any are available.
    pub(crate) fn pop_scope_layer(&mut self) -> Option<BTreeMap<Sym, Name>> {
        self.scope.pop()
    }

    /// Collects local variables for the given AST body non-recursively.
    pub(crate) fn collect_locals(&mut self, body: &Body) {
        locals::CollectLocals::new(self).collect(body);
    }
}

#[derive(Debug, Default)]
pub struct ConstData {
    ints: HashMap<i64, (ConstHandle, IntRef)>,
    strs: HashMap<String, (ConstHandle, StrRef)>,
    const_pool: ConstPool,
}

impl ConstData {
    /// Gets or inserts a static int reference.
    pub fn const_int(&mut self, int: i64) -> (ConstHandle, IntRef) {
        if let Some(pair) = self.ints.get(&int) {
            pair.clone()
        } else {
            let obj = Int::new_obj(int);
            let hdl = self.const_pool.push(obj.clone());
            self.ints.insert(int, (hdl, obj.clone()));
            (hdl, obj)
        }
    }

    /// Gets or inserts a static string reference.
    pub fn const_str(&mut self, s: impl AsRef<str>) -> (ConstHandle, StrRef) {
        let s = s.as_ref();
        if let Some(pair) = self.strs.get(s) {
            pair.clone()
        } else {
            let obj = Str::new_obj(s.to_string());
            let hdl = self.const_pool.push(obj.clone());
            self.strs.insert(s.to_string(), (hdl, obj.clone()));
            (hdl, obj)
        }
    }
}