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

use super::visitor::*;

pub(crate) struct PassRewriting {
    current_scope: BlockStatementId,
    statement_buffer: ScopedBuffer<StatementId>,
    call_expr_buffer: ScopedBuffer<CallExpressionId>,
    expression_buffer: ScopedBuffer<ExpressionId>,
}

impl PassRewriting {
    pub(crate) fn new() -> Self {
        Self{
            current_scope: BlockStatementId::new_invalid(),
            statement_buffer: ScopedBuffer::with_capacity(16),
            call_expr_buffer: ScopedBuffer::with_capacity(16),
            expression_buffer: ScopedBuffer::with_capacity(16),
        }
    }
}

impl Visitor for PassRewriting {
    // --- Visiting procedures

    fn visit_component_definition(&mut self, ctx: &mut Ctx, id: ComponentDefinitionId) -> VisitorResult {
        let def = &ctx.heap[id];
        let body_id = def.body;
        return self.visit_block_stmt(ctx, body_id);
    }

    fn visit_function_definition(&mut self, ctx: &mut Ctx, id: FunctionDefinitionId) -> VisitorResult {
        let def = &ctx.heap[id];
        let body_id = def.body;
        return self.visit_block_stmt(ctx, body_id);
    }

    // --- Visiting statements (that are not the select statement)

    fn visit_block_stmt(&mut self, ctx: &mut Ctx, id: BlockStatementId) -> VisitorResult {
        let block_stmt = &ctx.heap[id];
        let stmt_section = self.statement_buffer.start_section_initialized(&block_stmt.statements);

        self.current_scope = id;
        for stmt_idx in 0..stmt_section.len() {
            self.visit_stmt(ctx, stmt_section[stmt_idx])?;
        }

        stmt_section.forget();
        return Ok(())
    }

    fn visit_labeled_stmt(&mut self, ctx: &mut Ctx, id: LabeledStatementId) -> VisitorResult {
        let labeled_stmt = &ctx.heap[id];
        let body_id = labeled_stmt.body;
        return self.visit_stmt(ctx, body_id);
    }

    fn visit_if_stmt(&mut self, ctx: &mut Ctx, id: IfStatementId) -> VisitorResult {
        let if_stmt = &ctx.heap[id];
        let true_body_id = if_stmt.true_body;
        let false_body_id = if_stmt.false_body;

        self.visit_block_stmt(ctx, true_body_id)?;
        if let Some(false_body_id) = false_body_id {
            self.visit_block_stmt(ctx, false_body_id)?;
        }

        return Ok(())
    }

    fn visit_while_stmt(&mut self, ctx: &mut Ctx, id: WhileStatementId) -> VisitorResult {
        let while_stmt = &ctx.heap[id];
        let body_id = while_stmt.body;
        return self.visit_block_stmt(ctx, body_id);
    }

    fn visit_synchronous_stmt(&mut self, ctx: &mut Ctx, id: SynchronousStatementId) -> VisitorResult {
        let sync_stmt = &ctx.heap[id];
        let body_id = sync_stmt.body;
        return self.visit_block_stmt(ctx, body_id);
    }

    // --- Visiting the select statement

    fn visit_select_stmt(&mut self, ctx: &mut Ctx, id: SelectStatementId) -> VisitorResult {
        // We're going to transform the select statement by a block statement
        // containing builtin runtime-calls. And to do so we create temporary
        // variables and move some other statements around.
        let select_stmt = &ctx.heap[id];
        let mut total_num_cases = select_stmt.cases.len();
        let mut total_num_ports = 0;

        // Put heap IDs into temporary buffers to handle borrowing rules
        let mut call_id_section = self.call_expr_buffer.start_section();
        let mut expr_id_section = self.expression_buffer.start_section();

        for case in select_stmt.cases.iter() {
            total_num_ports += case.involved_ports.len();
            for (call_id, expr_id) in case.involved_ports.iter().copied() {
                call_id_section.push(call_id);
                expr_id_section.push(expr_id);
            }
        }

        // Transform all of the call expressions by takings its argument (the
        // port from which we `get`) and turning it into a temporary variable.
        let mut transformed_stmts = Vec::with_capacity(total_num_ports); // TODO: Recompute this preallocated length, put assert at the end
        let mut locals = Vec::with_capacity(total_num_ports);

        for port_var_idx in 0..call_id_section.len() {
            let call_id = call_id_section[port_var_idx];
            let expr_id = expr_id_section[port_var_idx];
            let (replacement_variable_id, variable_stmt_id) = self.modify_call_and_create_replacement_variable(ctx, call_id, expr_id);
            transformed_stmts.push(variable_stmt_id);
            locals.push(replacement_variable_id);
        }

        call_id_section.forget();
        expr_id_section.forget();

        // let block = ctx.heap.alloc_block_statement(|this| BlockStatement{
        //     this,
        //     is_implicit: true,
        //     span: stmt.span,
        //     statements: vec![],
        //     end_block: EndBlockStatementId(),
        //     scope_node: ScopeNode {},
        //     first_unique_id_in_scope: 0,
        //     next_unique_id_in_scope: 0,
        //     locals,
        //     labels: vec![],
        //     next: ()
        // });

        return Ok(())
    }
}

impl PassRewriting {
    fn modify_call_and_create_replacement_variable(&self, ctx: &mut Ctx, call_expr_id: CallExpressionId, port_expr_id: ExpressionId) -> (VariableId, MemoryStatementId) {
        // Retrieve original expression which we're going to transplant into
        // its own variable
        let port_expr = &ctx.heap[port_expr_id];
        let port_expr_span = port_expr.full_span();
        let port_expr_unique_id = port_expr.get_unique_id_in_definition();

        // Create the entries in the heap
        let variable_expr_id = ctx.heap.alloc_variable_expression(|this| VariableExpression{
            this,
            identifier: Identifier::new_empty(port_expr_span),
            declaration: None,
            used_as_binding_target: false,
            parent: ExpressionParent::None,
            unique_id_in_definition: port_expr_unique_id,
        });
        let initial_expr_id = ctx.heap.alloc_assignment_expression(|this| AssignmentExpression{
            this,
            operator_span: port_expr_span,
            full_span: port_expr_span,
            left: variable_expr_id.upcast(),
            operation: AssignmentOperator::Set,
            right: port_expr_id,
            parent: ExpressionParent::None,
            unique_id_in_definition: -1,
        });
        let variable_id = ctx.heap.alloc_variable(|this| Variable{
            this,
            kind: VariableKind::Local,
            parser_type: ParserType{ elements: vec![ParserTypeElement{
                    element_span: port_expr_span,
                    variant: ParserTypeVariant::Inferred,
                }],
                full_span: port_expr_span
            },
            identifier: Identifier::new_empty(port_expr_span),
            relative_pos_in_block: 0,
            unique_id_in_scope: 0,
        });
        let variable_decl_stmt = ctx.heap.alloc_memory_statement(|this| MemoryStatement{
            this,
            span: port_expr_span,
            variable: variable_id,
            initial_expr: initial_expr_id,
            next: StatementId::new_invalid(),
        });

        // Modify all entries that required access other heap entries
        let variable_expr = &mut ctx.heap[variable_expr_id];
        variable_expr.declaration = Some(variable_id);
        variable_expr.parent = ExpressionParent::Expression(initial_expr_id.upcast(), 1);

        let initial_expr = &mut ctx.heap[initial_expr_id];
        initial_expr.parent = ExpressionParent::Memory(variable_decl_stmt);

        // Modify the parent of the expression that we just transplanted
        let port_expr = &mut ctx.heap[port_expr_id];
        *port_expr.parent_mut() = ExpressionParent::Expression(initial_expr_id.upcast(), 1);

        // Modify the call expression (that should contain the port expression
        // as the first argument) to point to the new variable
        let call_arg_expr_id = ctx.heap.alloc_variable_expression(|this| VariableExpression{
            this,
            identifier: Identifier::new_empty(port_expr_span),
            declaration: Some(variable_id),
            used_as_binding_target: false,
            parent: ExpressionParent::Expression(call_expr_id.upcast(), 0),
            unique_id_in_definition: port_expr_unique_id,
        });
        let call_expr = &mut ctx.heap[call_expr_id];
        debug_assert!(call_expr.arguments.len() == 1 && call_expr.arguments[0] == port_expr_id);
        call_expr.arguments[0] = call_arg_expr_id.upcast();

        return (variable_id, variable_decl_stmt);
    }
}