// 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 get_call_expr_id = call_id_section[port_var_idx];
            let port_expr_id = expr_id_section[port_var_idx];

            // Move the port expression such that it gets assigned to a temporary variable
            let variable_id = create_ast_variable(ctx);
            let variable_decl_stmt_id = create_ast_variable_declaration_stmt(ctx, variable_id, port_expr_id);

            // Replace the original port expression in the call with a reference
            // to the replacement variable
            let variable_expr_id = create_ast_variable_expr(ctx, variable_id);
            let call_expr = &mut ctx.heap[get_call_expr_id];
            call_expr.arguments[0] = variable_expr_id.upcast();

            transformed_stmts.push(variable_decl_stmt_id.upcast().upcast());
            locals.push(variable_id);
        }

        // Insert runtime calls that facilitate the semantics of the select
        // block.

        // Create the call that indicates the start of the select block
        {
            let num_cases_expression_id = create_ast_literal_integer_expr(ctx, total_num_cases as u64);
            let num_ports_expression_id = create_ast_literal_integer_expr(ctx, total_num_ports as u64);
            let arguments = vec![
                num_cases_expression_id.upcast(),
                num_ports_expression_id.upcast()
            ];

            let call_expression_id = create_ast_call_expr(ctx, Method::SelectStart, arguments);
            let call_statement_id = create_ast_expression_stmt(ctx, call_expression_id.upcast());

            transformed_stmts.push(call_statement_id.upcast());
        }

        // Create calls for each select case that will register the ports that
        // we are waiting on at the runtime.
        {
            let mut total_port_index = 0;
            for case_index in 0..total_num_cases {
                let case = &ctx.heap[id].cases[case_index];
                let case_num_ports = case.involved_ports.len();

                for case_port_index in 0..case_num_ports {
                    // Arguments to runtime call
                    let port_variable_id = locals[total_port_index]; // so far this variable contains the temporary variables for the port expressions
                    let case_index_expr_id = create_ast_literal_integer_expr(ctx, case_index as u64);
                    let port_index_expr_id = create_ast_literal_integer_expr(ctx, case_port_index as u64);
                    let port_variable_expr_id = create_ast_variable_expr(ctx, port_variable_id);
                    let runtime_call_arguments = vec![
                        case_index_expr_id.upcast(),
                        port_index_expr_id.upcast(),
                        port_variable_expr_id.upcast()
                    ];

                    // Create runtime call, then store it
                    let runtime_call_expr_id = create_ast_call_expr(ctx, Method::SelectRegisterCasePort, runtime_call_arguments);
                    let runtime_call_stmt_id = create_ast_expression_stmt(ctx, runtime_call_expr_id);

                    transformed_stmts.push(runtime_call_stmt_id.upcast());

                    total_port_index += 1;
                }
            }
        }

        // Create the variable that will hold the result of a completed select
        // block. Then create the runtime call that will produce this result
        let select_variable_id = create_ast_variable(ctx);
        locals.push(select_variable_id);

        {
            let runtime_call_expr_id = create_ast_call_expr(ctx, Method::SelectWait, Vec::new());
            let variable_stmt_id = create_ast_variable_declaration_stmt(ctx, select_variable_id, runtime_call_expr_id.upcast());
            transformed_stmts.push(variable_stmt_id.upcast().upcast());
        }

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

        // Now we transform each of the select block case's guard and code into
        // a chained if-else statement.
        if total_num_cases > 0 {
            let cmp_expr_id = create_ast_equality_comparison_expr(ctx, select_variable_id, 0);
            let mut (if_stmt_id, end_if_stmt_id) = create_ast_if_statement(ctx, cmp_expr_id.upcast(), )
        }

        // 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 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_wait_variable_and_statement(&self, ctx: &mut Ctx) -> (VariableId, MemoryStatementId) {
        let variable_id = create_ast_variable(ctx);
        let variable_expr_id = create_ast_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);
    }
}

// -----------------------------------------------------------------------------
// Utilities to create compiler-generated AST nodes
// -----------------------------------------------------------------------------

fn create_ast_variable(ctx: &mut Ctx) -> VariableId {
    return 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,
    });
}

fn create_ast_variable_expr(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
    });
}

fn create_ast_call_expr(ctx: &mut Ctx, method: Method, arguments: Vec<ExpressionId>) -> CallExpressionId {
    let call_expression_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,
    });

    for (argument_index, argument_id) in arguments.iter().cloned().enumerate() {
        let argument_expr = &mut ctx.heap[argument_id];
        *argument_expr.parent_mut() = ExpressionParent::Expression(call_expression_id.upcat(), argument_index as u32);
    }

    return call_expression_id;
}

fn create_ast_literal_integer_expr(ctx: &mut Ctx, unsigned_value: u64) -> LiteralExpressionId {
    return ctx.heap.alloc_literal_expression(|this| LiteralExpression{
        this,
        span: InputSpan::new(),
        value: Literal::Integer(LiteralInteger{
            unsigned_value,
            negated: false,
        }),
        parent: ExpressionParent::None,
        unique_id_in_definition: -1
    });
}

fn create_ast_equality_comparison_expr(ctx: &mut Ctx, variable_id: VariableId, value: u64) -> BinaryExpressionId {
    let var_expr_id = create_ast_variable_expr(ctx, variable_id);
    let int_expr_id = create_ast_literal_integer_expr(ctx, value);
    let cmp_expr_id = ctx.heap.alloc_binary_expression(|this| BinaryExpression{
        this,
        operator_span: InputSpan::new(),
        full_span: InputSpan::new(),
        left: var_expr_id.upcast(),
        operation: BinaryOperator::Equality,
        right: int_expr_id.upcast(),
        parent: ExpressionParent::None,
        unique_id_in_definition: -1,
    });

    let var_expr = &mut ctx.heap[var_expr_id];
    var_expr.parent = ExpressionParent::Expression(cmp_expr_id.upcast(), 0);
    let int_expr = &mut ctx.heap[int_expr_id];
    int_expr.parent = ExpressionParent::Expression(cmp_expr_id.upcast(), 1);

    return cmp_expr_id;
}

fn create_ast_expression_stmt(ctx: &mut Ctx, expression_id: ExpressionId) -> ExpressionStatementId {
    let statement_id = ctx.heap.alloc_expression_statement(|this| ExpressionStatement{
        this,
        span: InputSpan::new(),
        expression: expression_id,
        next: StatementId::new_invalid(),
    });

    let expression = &mut ctx.heap[expression_id];
    *expression.parent_mut() = ExpressionParent::ExpressionStmt(statement_id);

    return statement_id;
}

fn create_ast_variable_declaration_stmt(ctx: &mut Ctx, variable_id: VariableId, initial_value_expr_id: ExpressionId) -> MemoryStatementId {
    // Create the assignment expression, assigning the initial value to the variable
    let variable_expr_id = create_ast_variable_expr(ctx, variable_id);
    let assignment_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: initial_value_expr_id,
        parent: ExpressionParent::None,
        unique_id_in_definition: -1,
    });

    // Create the memory statement
    let memory_stmt_id = ctx.heap.alloc_memory_statement(|this| MemoryStatement{
        this,
        span: InputSpan::new(),
        variable: variable_id,
        initial_expr: assignment_expr_id,
        next: StatementId::new_invalid(),
    });

    // Set all parents which we can access
    let variable_expr = &mut ctx.heap[variable_expr_id];
    variable_expr.parent = ExpressionParent::Expression(assignment_expr_id.upcast(), 0);
    let value_expr = &mut ctx.heap[initial_value_expr_id];
    *value_expr.parent_mut() = ExpressionParent::Expression(assignment_expr_id.upcast(), 1);
    let assignment_expr = &mut ctx.heap[assignment_expr_id];
    assignment_expr.parent = ExpressionParent::Memory(memory_stmt_id);

    return memory_stmt_id;
}

fn create_ast_if_statement(ctx: &mut Ctx, condition_expression_id: ExpressionId, true_body: BlockStatementId, false_body: Option<BlockStatementId>) -> (IfStatementId, EndIfStatementId) {
    // Create if statement and the end-if statement
    let if_stmt_id = ctx.heap.alloc_if_statement(|this| IfStatement{
        this,
        span: InputSpan::new(),
        test: condition_expression_id,
        true_body,
        false_body,
        end_if: EndIfStatementId::new_invalid()
    });

    let end_if_stmt_id = ctx.heap.alloc_end_if_statement(|this| EndIfStatement{
        this,
        start_if: if_stmt_id,
        next: StatementId::new_invalid(),
    });

    // Link the statements up as much as we can
    let if_stmt = &mut ctx.heap[if_stmt_id];
    if_stmt.end_if = end_if_stmt_id;

    let condition_expr = &mut ctx.heap[condition_expression_id];
    *condition_expr.parent_mut() = ExpressionParent::If(if_stmt_id);

    let true_body_stmt = &ctx.heap[true_body];
    let true_body_end_stmt = &mut ctx.heap[true_body_stmt.end_block];
    true_body_end_stmt.next = end_if_stmt_id.upcast();

    if let Some(false_body) = false_body {
        let false_body_stmt = &ctx.heap[false_body];
        let false_body_end_stmt = &mut ctx.heap[false_body_stmt.end_block];
        false_body_end_stmt.next = end_if_stmt_id.upcast();
    }

    return (if_stmt_id, end_if_stmt_id);
}

fn set_ast_if_statement_false_body(ctx: &mut Ctx, if_statement_id: IfStatementId, end_if_statement_id: EndIfStatementId, false_body: BlockStatementId) {
    // Point if-statement to "false body"
    let if_stmt = &mut ctx.heap[if_statement_id];
    debug_assert!(if_stmt.false_body.is_none()); // simplifies logic, not necessary
    if_stmt.false_body = Some(false_body);

    // Point end of false body to the end of the if statement
    let false_body_stmt = &ctx.heap[false_body];
    let false_body_end_stmt = &mut ctx.heap[false_body_stmt.end_block];
    false_body_end_stmt.next = end_if_statement_id.upcast();
}