use crate::ir::ir_block::IrBlock;
use crate::ir::ir_parameter::IrParameter;
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_variable::IrVariable;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

pub struct IrBuilder {
    parameters: Vec<Rc<IrParameter>>,
    block_counter: usize,
    t_var_counter: usize,
    blocks: HashMap<usize, Rc<RefCell<IrBlock>>>,
    current_block_builder: Option<IrBlockBuilder>,
}

impl IrBuilder {
    pub fn new() -> Self {
        Self {
            parameters: vec![],
            block_counter: 0,
            t_var_counter: 0,
            blocks: HashMap::new(),
            current_block_builder: None,
        }
    }

    pub fn parameters(&self) -> &[Rc<IrParameter>] {
        &self.parameters
    }

    pub fn parameters_mut(&mut self) -> &mut Vec<Rc<IrParameter>> {
        &mut self.parameters
    }

    pub fn new_block(&mut self) -> usize {
        let block_id = self.block_counter;
        self.block_counter += 1;
        let block_builder = IrBlockBuilder::new(block_id);
        self.current_block_builder = Some(block_builder);
        block_id
    }

    pub fn get_block(&mut self, block_id: usize) -> &Rc<RefCell<IrBlock>> {
        self.blocks
            .get(&block_id)
            .expect(&format!("Block {} not found", block_id))
    }

    pub fn current_block_mut(&mut self) -> &mut IrBlockBuilder {
        self.current_block_builder
            .as_mut()
            .expect("No current block builder")
    }

    pub fn finish_block(&mut self) {
        let builder = self
            .current_block_builder
            .take()
            .expect("No current block builder");
        let block = builder.build();
        self.blocks.insert(block.id(), Rc::new(RefCell::new(block)));
    }

    pub fn new_t_var(&mut self) -> String {
        let id = self.t_var_counter;
        self.t_var_counter += 1;
        format!("t{}", id)
    }
}

pub struct IrBlockBuilder {
    id: usize,
    statements: Vec<IrStatement>,
}

impl IrBlockBuilder {
    pub fn new(id: usize) -> Self {
        Self {
            id,
            statements: vec![],
        }
    }

    pub fn add_statement(&mut self, statement: IrStatement) {
        self.statements.push(statement);
    }

    pub fn build(self) -> IrBlock {
        IrBlock::new(self.id, &format!("b{}", self.id), self.statements)
    }
}