// TODO: File contains a lot of manual AST element construction. Wherein we have
//  (for the first time in this compiler) a lot of fields that have no real
//  meaning (e.g. the InputSpan of a AST-transformation). What are we going to
//  do with this to make the code and datastructures more easily grokable?
//  We could do an intermediate AST structure. But considering how close this
//  phase of compilation is to bytecode generation, that might be a lot of busy-
//  work with few results. Alternatively we may put the AST elements inside
//  a special substructure. We could also force ourselves (and put the
//  appropriate comments in the code) to not use certain fields anymore after
//  a particular stage of compilation.

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_get_call_insert_variable(ctx, call_id, expr_id);
            transformed_stmts.push(variable_stmt_id.upcast().upcast());
            locals.push(replacement_variable_id);
        }

        // Our transformed statements now contain all of the temporary port
        // calculations. We'll now insert the appropriate runtime calls to
        // notify the runtime that we're going to wait on these ports.
        let (_, select_start_stmt_id) = self.create_runtime_select_start_call_statement(ctx, total_num_cases, total_num_ports);
        transformed_stmts.push(select_start_stmt_id.upcast());

        // TODO: Call the runtime function for eeach of the substituted port variables to register all ports for the select statement

        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_get_call_insert_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: -1,
            unique_id_in_scope: -1,
        });
        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_eq!(call_expr.method, Method::Get);
        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);
    }

    fn create_runtime_call_statement(&self, ctx: &mut Ctx, method: Method, arguments: Vec<ExpressionId>) -> (CallExpressionId, ExpressionStatementId) {
        let call_expr_id = ctx.heap.alloc_call_expression(|this| CallExpression{
            this,
            func_span: InputSpan::new(),
            full_span: InputSpan::new(),
            parser_type: ParserType{
                elements: Vec::new(),
                full_span: InputSpan::new(),
            },
            method,
            arguments,
            definition: DefinitionId::new_invalid(),
            parent: ExpressionParent::None,
            unique_id_in_definition: -1,
        });
        let call_stmt_id = ctx.heap.alloc_expression_statement(|this| ExpressionStatement{
            this,
            span: InputSpan::new(),
            expression: call_expr_id.upcast(),
            next: StatementId::new_invalid(),
        });

        let call_expr = &mut ctx.heap[call_expr_id];
        call_expr.parent = ExpressionParent::ExpressionStmt(call_stmt_id);

        return (call_expr_id, call_stmt_id);
    }

    fn create_runtime_select_start_call_statement(&self, ctx: &mut Ctx, num_cases: usize, num_ports_total: usize) -> (CallExpressionId, ExpressionStatementId) {
        let num_cases_expr_id = self.create_literal_integer(ctx, num_cases as u64);
        let num_ports_expr_id = self.create_literal_integer(ctx, num_ports_total as u64);
        let arguments = vec![
            num_cases_expr_id.upcast(),
            num_ports_expr_id.upcast()
        ];

        let (call_expr_id, call_stmt_id) = self.create_runtime_call_statement(ctx, Method::SelectStart, arguments);

        let num_cases_expr = &mut ctx.heap[num_cases_expr_id];
        num_cases_expr.parent = ExpressionParent::Expression(call_expr_id.upcast(), 0);
        let num_ports_expr = &mut ctx.heap[num_ports_expr_id];
        num_ports_expr.parent = ExpressionParent::Expression(num_ports_expr_id.upcast(), 1);

        return (call_expr_id, call_stmt_id);
    }

    fn create_runtime_select_register_port_call_statement(&self, ctx: &mut Ctx, case_index: usize, port_index: usize, original_port_expr_id: ExpressionId, port_variable_id: VariableId) -> (CallExpressionId, ExpressionStatementId) {
        let original_port_expr = &ctx.heap[original_port_expr_id];
        let original_port_span = original_port_expr.full_span();
        let original_port_unique_id = original_port_expr.get_unique_id_in_definition();

        let case_index_expr_id = self.create_literal_integer(ctx, case_index as u64);
        let port_index_expr_id = self.create_literal_integer(ctx, port_index as u64);
        let port_var_expr_id = self.create_variable_expr(ctx, port_variable_id);

        let arguments = vec![
            case_index_expr_id.upcast(),
            port_index_expr_id.upcast(),
            port_var_expr_id.upcast()
        ];

        let (call_expr_id, call_stmt_id) = self.create_runtime_call_statement(ctx, Method::SelectRegisterCasePort, arguments);

        let case_index_expr = &mut ctx.heap[case_index_expr_id];
        case_index_expr.parent = ExpressionParent::Expression(call_expr_id.upcast(), 0);
        let port_index_expr = &mut ctx.heap[port_index_expr_id];
        port_index_expr.parent = ExpressionParent::Expression(call_expr_id.upcast(), 1);
        let port_var_expr = &mut ctx.heap[port_var_expr_id];
        port_var_expr.parent = ExpressionParent::Expression(call_expr_id.upcast(), 2);

        return (call_expr_id, call_stmt_id);
    }

    fn create_runtime_select_wait_variable_and_statement(&self, ctx: &mut Ctx) -> (VariableId, MemoryStatementId) {
        let variable_id = ctx.heap.alloc_variable(|this| Variable{
            this,
            kind: VariableKind::Local,
            parser_type: ParserType{
                elements: Vec::new(),
                full_span: InputSpan::new(),
            },
            identifier: Identifier::new_empty(InputSpan::new()),
            relative_pos_in_block: -1,
            unique_id_in_scope: -1
        });
        let variable_expr_id = self.create_variable_expr(ctx, variable_id);
        let runtime_call_expr_id = ctx.heap.alloc_call_expression(|this| CallExpression{
            this,
            func_span: InputSpan::new(),
            full_span: InputSpan::new(),
            parser_type: ParserType{
                elements: Vec::new(),
                full_span: InputSpan::new(),
            },
            method: Method::SelectWait,
            arguments: Vec::new(),
            definition: DefinitionId::new_invalid(),
            parent: ExpressionParent::None,
            unique_id_in_definition: -1
        });
        let initial_expr_id = ctx.heap.alloc_assignment_expression(|this| AssignmentExpression{
            this,
            operator_span: InputSpan::new(),
            full_span: InputSpan::new(),
            left: variable_expr_id.upcast(),
            operation: AssignmentOperator::Set,
            right: runtime_call_expr_id.upcast(),
            parent: ExpressionParent::None,
            unique_id_in_definition: -1
        });

        let variable_statement_id = ctx.heap.alloc_memory_statement(|this| MemoryStatement{
            this,
            span: InputSpan::new(),
            variable: variable_id,
            initial_expr: initial_expr_id,
            next: StatementId::new_invalid()
        });

        let variable_expr = &mut ctx.heap[variable_expr_id];
        variable_expr.parent = ExpressionParent::Expression(initial_expr_id.upcast(), 0);
        let runtime_call_expr = &mut ctx.heap[runtime_call_expr_id];
        runtime_call_expr.parent = ExpressionParent::Expression(initial_expr_id.upcast(), 1);
        let initial_expr = &mut ctx.heap[initial_expr_id];
        initial_expr.parent = ExpressionParent::Memory(variable_statement_id);

        return (variable_id, variable_statement_id);
    }

    /// Creates an integer literal. The caller still needs to set its expression
    /// parent afterwards.
    fn create_literal_integer(&self, ctx: &mut Ctx, value: u64) -> LiteralExpressionId {
        return ctx.heap.alloc_literal_expression(|this| LiteralExpression{
            this,
            span: InputSpan::new(),
            value: Literal::Integer(LiteralInteger{
                unsigned_value: value,
                negated: false,
            }),
            parent: ExpressionParent::None,
            unique_id_in_definition: -1
        });
    }

    fn create_variable_expr(&self, ctx: &mut Ctx, variable_id: VariableId) -> VariableExpressionId {
        return ctx.heap.alloc_variable_expression(|this| VariableExpression{
            this,
            identifier: Identifier::new_empty(InputSpan::new()),
            declaration: Some(variable_id),
            used_as_binding_target: false,
            parent: ExpressionParent::None,
            unique_id_in_definition: -1
        })
    }
}