//! Control flow structures for the JIT compiler.

use crate::loader::Instruction;
use crate::loader::decode::{BifOp, MapOp};
use cranelift_frontend::FunctionBuilder;
use std::collections::{HashMap, HashSet};

use super::compiler::JitError;
use super::ir_arithmetic::{ArithmeticOp, ParsedBif};
use super::ir_common::{
    ensure_known_label, validate_label_operand, validate_read_operand, validate_write_operand,
};
use super::ir_control_validation::{
    validate_float_register_operand, validate_fmove_operands, validate_import_operand,
    validate_supported_type_test,
};
use super::ir_guards::{
    immediate_raw_term, immediate_usize, parse_select_pairs, validate_tag_atom,
};
use super::ir_map::{
    parse_get_map_elements_operands, parse_has_map_fields_operands, parse_put_map_operands,
};

pub(crate) struct TranslationPlan {
    pub(crate) labels: HashMap<u32, usize>,
    pub(crate) block_starts: HashSet<usize>,
}

impl TranslationPlan {
    pub(crate) fn new(instructions: &[Instruction]) -> Result<Self, JitError> {
        if instructions.is_empty() {
            return Err(JitError::EmptyFunction);
        }

        let mut labels = HashMap::new();
        let mut block_starts = HashSet::from([0, instructions.len()]);
        for (index, instruction) in instructions.iter().enumerate() {
            match instruction {
                Instruction::Label { label } => {
                    labels.insert(*label, index);
                    block_starts.insert(index);
                }
                Instruction::Return => {
                    block_starts.insert(index + 1);
                }
                Instruction::Move {
                    source,
                    destination,
                } => {
                    validate_read_operand(source)?;
                    validate_write_operand(destination)?;
                }
                Instruction::Swap { left, right } => {
                    validate_read_operand(left)?;
                    validate_read_operand(right)?;
                    validate_write_operand(left)?;
                    validate_write_operand(right)?;
                }
                Instruction::TypeTest { op, fail, value } => {
                    validate_supported_type_test(*op)?;
                    validate_label_operand(fail)?;
                    validate_read_operand(value)?;
                    block_starts.insert(index + 1);
                }
                Instruction::PutList {
                    head,
                    tail,
                    destination,
                } => {
                    validate_read_operand(head)?;
                    validate_read_operand(tail)?;
                    validate_write_operand(destination)?;
                    block_starts.insert(index + 1);
                }
                Instruction::GetList { source, head, tail } => {
                    validate_read_operand(source)?;
                    validate_write_operand(head)?;
                    validate_write_operand(tail)?;
                }
                Instruction::GetHd {
                    source,
                    destination,
                }
                | Instruction::GetTl {
                    source,
                    destination,
                } => {
                    validate_read_operand(source)?;
                    validate_write_operand(destination)?;
                }
                Instruction::PutTuple2 {
                    destination,
                    elements,
                } => {
                    validate_write_operand(destination)?;
                    let crate::loader::decode::Operand::List(elements) = elements else {
                        return Err(JitError::UnsupportedOperand {
                            operand: format!(
                                "put_tuple2 elements must be a list, got {elements:?}"
                            ),
                        });
                    };
                    for element in elements {
                        validate_read_operand(element)?;
                    }
                    block_starts.insert(index + 1);
                }
                Instruction::GetTupleElement {
                    source,
                    index,
                    destination,
                } => {
                    validate_read_operand(source)?;
                    let _ = immediate_usize(index, "get_tuple_element index")?;
                    validate_write_operand(destination)?;
                }
                Instruction::Comparison {
                    fail, left, right, ..
                } => {
                    validate_label_operand(fail)?;
                    validate_read_operand(left)?;
                    validate_read_operand(right)?;
                    block_starts.insert(index + 1);
                }
                Instruction::TestArity { fail, tuple, arity } => {
                    validate_label_operand(fail)?;
                    validate_read_operand(tuple)?;
                    let _ = immediate_usize(arity, "test_arity arity")?;
                    block_starts.insert(index + 1);
                }
                Instruction::IsTaggedTuple {
                    fail,
                    value,
                    arity,
                    tag,
                } => {
                    validate_label_operand(fail)?;
                    validate_read_operand(value)?;
                    let _ = immediate_usize(arity, "is_tagged_tuple arity")?;
                    validate_tag_atom(tag)?;
                    block_starts.insert(index + 1);
                }
                Instruction::SelectVal { value, fail, list } => {
                    validate_read_operand(value)?;
                    validate_label_operand(fail)?;
                    for (candidate, target) in parse_select_pairs(list)? {
                        let _ = immediate_raw_term(candidate)?;
                        validate_label_operand(target)?;
                    }
                    block_starts.insert(index + 1);
                }
                Instruction::Jump { target } => {
                    validate_label_operand(target)?;
                    block_starts.insert(index + 1);
                }
                Instruction::Try { destination, label } => {
                    validate_write_operand(destination)?;
                    validate_label_operand(label)?;
                    block_starts.insert(index + 1);
                }
                Instruction::TryEnd { source } | Instruction::TryCase { source } => {
                    validate_write_operand(source)?;
                    block_starts.insert(index + 1);
                }
                Instruction::Call { label, .. } | Instruction::CallOnly { label, .. } => {
                    validate_label_operand(label)?;
                    block_starts.insert(index + 1);
                }
                Instruction::CallExt { import, .. } | Instruction::CallExtOnly { import, .. } => {
                    validate_import_operand(import)?;
                    block_starts.insert(index + 1);
                }
                Instruction::Bif { op, operands } => {
                    let parsed = ParsedBif::parse(*op, operands)?;
                    let _ = ArithmeticOp::from_import(parsed.import)?;
                    validate_label_operand(parsed.fail)?;
                    validate_read_operand(parsed.left)?;
                    validate_read_operand(parsed.right)?;
                    validate_write_operand(parsed.destination)?;
                    block_starts.insert(index + 1);
                }
                Instruction::MakeFun { .. } => {
                    block_starts.insert(index + 1);
                }
                Instruction::Fmove { source, dest } => {
                    validate_fmove_operands(source, dest)?;
                    block_starts.insert(index + 1);
                }
                Instruction::Fconv { source, dest } => {
                    validate_read_operand(source)?;
                    validate_float_register_operand(dest, "fconv destination")?;
                    block_starts.insert(index + 1);
                }
                Instruction::Fadd {
                    fail,
                    left,
                    right,
                    dest,
                }
                | Instruction::Fsub {
                    fail,
                    left,
                    right,
                    dest,
                }
                | Instruction::Fmul {
                    fail,
                    left,
                    right,
                    dest,
                }
                | Instruction::Fdiv {
                    fail,
                    left,
                    right,
                    dest,
                } => {
                    validate_label_operand(fail)?;
                    validate_float_register_operand(left, "float arithmetic left")?;
                    validate_float_register_operand(right, "float arithmetic right")?;
                    validate_float_register_operand(dest, "float arithmetic destination")?;
                    block_starts.insert(index + 1);
                }
                Instruction::Fnegate { fail, source, dest } => {
                    validate_label_operand(fail)?;
                    validate_float_register_operand(source, "fnegate source")?;
                    validate_float_register_operand(dest, "fnegate destination")?;
                    block_starts.insert(index + 1);
                }
                Instruction::Apply { .. } | Instruction::CallFun { .. } => {
                    block_starts.insert(index + 1);
                }
                Instruction::BinaryOp { .. } => {
                    block_starts.insert(index + 1);
                }
                Instruction::MapOp { op, operands } => {
                    match op {
                        MapOp::PutMapAssoc | MapOp::PutMapExact => {
                            let (fail, parsed) = parse_put_map_operands(operands)?;
                            validate_label_operand(fail)?;
                            validate_read_operand(parsed.source)?;
                            validate_write_operand(parsed.destination)?;
                            for operand in parsed.pairs {
                                validate_read_operand(operand)?;
                            }
                        }
                        MapOp::GetMapElements => {
                            let (fail, parsed) = parse_get_map_elements_operands(operands)?;
                            validate_label_operand(fail)?;
                            validate_read_operand(parsed.source)?;
                            for pair in parsed.pairs.chunks_exact(2) {
                                validate_read_operand(&pair[0])?;
                                validate_write_operand(&pair[1])?;
                            }
                        }
                        MapOp::HasMapFields => {
                            let (fail, parsed) = parse_has_map_fields_operands(operands)?;
                            validate_label_operand(fail)?;
                            validate_read_operand(parsed.source)?;
                            for key in parsed.keys {
                                validate_read_operand(key)?;
                            }
                        }
                    }
                    block_starts.insert(index + 1);
                }
                Instruction::Send | Instruction::RemoveMessage | Instruction::Timeout => {
                    block_starts.insert(index + 1);
                }
                Instruction::LoopRec { fail, destination } => {
                    validate_label_operand(fail)?;
                    validate_write_operand(destination)?;
                    block_starts.insert(index + 1);
                }
                Instruction::LoopRecEnd { fail } | Instruction::Wait { fail } => {
                    validate_label_operand(fail)?;
                    block_starts.insert(index + 1);
                }
                Instruction::WaitTimeout { fail, timeout } => {
                    validate_label_operand(fail)?;
                    validate_read_operand(timeout)?;
                    block_starts.insert(index + 1);
                }
                Instruction::RecvMarkerReserve { dest } => {
                    validate_write_operand(dest)?;
                    block_starts.insert(index + 1);
                }
                Instruction::RecvMarkerBind { marker, label } => {
                    validate_read_operand(marker)?;
                    validate_label_operand(label)?;
                    block_starts.insert(index + 1);
                }
                Instruction::RecvMarkerClear { marker } | Instruction::RecvMarkerUse { marker } => {
                    validate_read_operand(marker)?;
                    block_starts.insert(index + 1);
                }
                other => {
                    return Err(JitError::UnsupportedOpcode {
                        opcode: opcode_name(other),
                    });
                }
            }
        }

        for instruction in instructions {
            match instruction {
                Instruction::TypeTest { fail, .. }
                | Instruction::Comparison { fail, .. }
                | Instruction::TestArity { fail, .. }
                | Instruction::IsTaggedTuple { fail, .. } => ensure_known_label(&labels, fail)?,
                Instruction::SelectVal { fail, list, .. } => {
                    ensure_known_label(&labels, fail)?;
                    for (_, target) in parse_select_pairs(list)? {
                        ensure_known_label(&labels, target)?;
                    }
                }
                Instruction::Jump { target }
                | Instruction::Try { label: target, .. }
                | Instruction::Call { label: target, .. }
                | Instruction::CallOnly { label: target, .. } => {
                    ensure_known_label(&labels, target)?
                }
                Instruction::Bif { op, operands } => {
                    if matches!(op, BifOp::Bif2 | BifOp::GcBif2) {
                        let parsed = ParsedBif::parse(*op, operands)?;
                        ensure_known_label(&labels, parsed.fail)?;
                    }
                }
                Instruction::Fadd { fail, .. }
                | Instruction::Fsub { fail, .. }
                | Instruction::Fmul { fail, .. }
                | Instruction::Fdiv { fail, .. }
                | Instruction::Fnegate { fail, .. } => {
                    ensure_known_label(&labels, fail)?;
                }
                Instruction::LoopRec { fail, .. }
                | Instruction::LoopRecEnd { fail }
                | Instruction::Wait { fail }
                | Instruction::WaitTimeout { fail, .. }
                | Instruction::RecvMarkerBind { label: fail, .. } => {
                    ensure_known_label(&labels, fail)?;
                }
                Instruction::MapOp { op, operands } => {
                    let fail = match op {
                        MapOp::PutMapAssoc | MapOp::PutMapExact => {
                            parse_put_map_operands(operands)?.0
                        }
                        MapOp::GetMapElements => parse_get_map_elements_operands(operands)?.0,
                        MapOp::HasMapFields => parse_has_map_fields_operands(operands)?.0,
                    };
                    ensure_known_label(&labels, fail)?;
                }
                _ => {}
            }
        }

        Ok(Self {
            labels,
            block_starts,
        })
    }
}

pub(crate) struct BlockMap {
    blocks_by_index: Vec<cranelift_codegen::ir::Block>,
    label_blocks: HashMap<u32, cranelift_codegen::ir::Block>,
    pub(crate) entry: cranelift_codegen::ir::Block,
    pub(crate) deopt: cranelift_codegen::ir::Block,
    pub(crate) exception_block: cranelift_codegen::ir::Block,
    pub(crate) yield_block: cranelift_codegen::ir::Block,
}

impl BlockMap {
    pub(crate) fn new(
        builder: &mut FunctionBuilder<'_>,
        instructions: &[Instruction],
        plan: &TranslationPlan,
    ) -> Self {
        let mut blocks_by_index = Vec::with_capacity(instructions.len() + 1);
        let entry = builder.create_block();
        builder.append_block_params_for_function_params(entry);
        let mut current = builder.create_block();
        for index in 0..=instructions.len() {
            if index > 0 && plan.block_starts.contains(&index) {
                current = builder.create_block();
            }
            blocks_by_index.push(current);
        }

        let mut label_blocks = HashMap::new();
        for (label, index) in &plan.labels {
            label_blocks.insert(*label, blocks_by_index[*index]);
        }

        Self {
            entry,
            blocks_by_index,
            label_blocks,
            deopt: builder.create_block(),
            exception_block: builder.create_block(),
            yield_block: builder.create_block(),
        }
    }

    pub(crate) fn block_for_instruction(&self, index: usize) -> cranelift_codegen::ir::Block {
        self.blocks_by_index[index]
    }

    pub(crate) fn block_after(&self, index: usize) -> cranelift_codegen::ir::Block {
        self.blocks_by_index[index + 1]
    }

    pub(crate) fn exit_block(&self) -> cranelift_codegen::ir::Block {
        self.blocks_by_index[self.blocks_by_index.len() - 1]
    }

    pub(crate) fn label_block(&self, label: u32) -> Result<cranelift_codegen::ir::Block, JitError> {
        self.label_blocks
            .get(&label)
            .copied()
            .ok_or(JitError::UnknownLabel { label })
    }
}

pub(crate) fn opcode_name(instruction: &Instruction) -> String {
    match instruction {
        Instruction::Generic { opcode, name, .. } => format!("{name} ({opcode})"),
        other => format!("{other:?}"),
    }
}
