use crate::collections::*;
use crate::protocol::*;

use super::visitor::*;

// Will get a rename. Will probably become bytecode emitter or something. For
// now it just scans the scopes and assigns a unique number for each variable
// such that, at any point in the program's execution, all accessible in-scope
// variables will have a unique position "on the stack".
pub(crate) struct PassStackSize {
    definition_buffer: ScopedBuffer<DefinitionId>,
    variable_buffer: ScopedBuffer<VariableId>,
    scope_buffer: ScopedBuffer<ScopeId>,
}

impl PassStackSize {
    pub(crate) fn new() -> Self {
        return Self{
            definition_buffer: ScopedBuffer::with_capacity(BUFFER_INIT_CAP_LARGE),
            variable_buffer: ScopedBuffer::with_capacity(BUFFER_INIT_CAP_SMALL),
            scope_buffer: ScopedBuffer::with_capacity(BUFFER_INIT_CAP_SMALL),
        }
    }
}

impl Visitor for PassStackSize {
    // Top level visitors

    fn visit_module(&mut self, ctx: &mut Ctx) -> VisitorResult {
        let module = ctx.module();
        debug_assert_eq!(module.phase, ModuleCompilationPhase::Rewritten);

        let root_id = module.root_id;
        let root = &ctx.heap[root_id];
        let definition_section = self.definition_buffer.start_section_initialized(&root.definitions);
        for definition_index in 0..definition_section.len() {
            let definition_id = definition_section[definition_index];
            self.visit_definition(ctx, definition_id)?
        }

        definition_section.forget();
        // ctx.module_mut().phase = ModuleCompilationPhase::StackSizeStuffAndStuff;
        return Ok(())
    }

    fn visit_function_definition(&mut self, ctx: &mut Ctx, id: FunctionDefinitionId) -> VisitorResult {
        let definition = &ctx.heap[id];
        let scope_id = definition.scope;

        self.visit_scope_and_assign_local_ids(ctx, scope_id, 0);
        return Ok(())
    }

    fn visit_component_definition(&mut self, ctx: &mut Ctx, id: ComponentDefinitionId) -> VisitorResult {
        let definition = &ctx.heap[id];
        let scope_id = definition.scope;

        self.visit_scope_and_assign_local_ids(ctx, scope_id, 0);
        return Ok(())
    }
}

impl PassStackSize {
    fn visit_scope_and_assign_local_ids(&mut self, ctx: &mut Ctx, scope_id: ScopeId, mut variable_counter: i32) {
        let scope = &mut ctx.heap[scope_id];
        scope.first_unique_id_in_scope = variable_counter;

        let variable_section = self.variable_buffer.start_section_initialized(&scope.variables);
        let child_scope_section = self.scope_buffer.start_section_initialized(&scope.nested);

        let mut variable_index = 0;
        let mut child_scope_index = 0;

        loop {
            // Determine relative positions of variable and scope to determine
            // which one occurs first within the current scope.
            let variable_relative_pos;
            if variable_index < variable_section.len() {
                let variable_id = variable_section[variable_index];
                let variable = &ctx.heap[variable_id];
                variable_relative_pos = variable.relative_pos_in_parent;
            } else {
                variable_relative_pos = i32::MAX;
            }

            let child_scope_relative_pos;
            if child_scope_index < child_scope_section.len() {
                let child_scope_id = child_scope_section[child_scope_index];
                let child_scope = &ctx.heap[child_scope_id];
                child_scope_relative_pos = child_scope.relative_pos_in_parent;
            } else {
                child_scope_relative_pos = i32::MAX;
            }

            if variable_relative_pos == i32::MAX && child_scope_relative_pos == i32::MAX {
                // Done, no more elements in the scope to consider
                break;
            }

            // Label the variable/scope, whichever comes first.
            if variable_relative_pos <= child_scope_relative_pos {
                debug_assert_ne!(variable_relative_pos, child_scope_relative_pos, "checking if this ever happens");
                let variable = &mut ctx.heap[variable_section[variable_index]];
                variable.unique_id_in_scope = variable_counter;
                variable_counter += 1;
                variable_index += 1;
            } else {
                let child_scope_id = child_scope_section[child_scope_index];
                self.visit_scope_and_assign_local_ids(ctx, child_scope_id, variable_counter);
                child_scope_index += 1;
            }
        }

        variable_section.forget();
        child_scope_section.forget();

        let scope = &mut ctx.heap[scope_id];
        scope.next_unique_id_in_scope = variable_counter;
    }
}