//! A frontend-oriented SSA builder for Slonik IR.

use std::collections::HashMap;

use crate::ir::{
    BinaryOp, Block, Body, CastOp, FloatCC, IntCC, MemSize, Stmt, StmtData, Type, UnaryOp, Value,
    Variable,
};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum IncomingEdgeKind {
    Jump,
    Then,
    Else,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct IncomingEdge {
    pred: Block,
    stmt: Stmt,
    kind: IncomingEdgeKind,
}

#[derive(Clone, Debug, Default)]
struct BlockState {
    /// Whether all predecessors of this block are known.
    sealed: bool,

    /// Whether a terminator has already been inserted into this block.
    filled: bool,

    /// Incoming CFG edges that target this block.
    preds: Vec<IncomingEdge>,

    /// Current SSA definitions visible at the end of this block.
    defs: HashMap<Variable, Value>,

    /// The variables that have synthesized block parameters in this block,
    /// in block-parameter order.
    param_order: Vec<Variable>,

    /// Mapping from variable to its synthesized block parameter value.
    params: HashMap<Variable, Value>,
}

/// Reusable scratch state for [`FrontendBuilder`].
#[derive(Default)]
pub struct FrontendBuilderContext {
    vars: HashMap<Variable, Type>,
    blocks: HashMap<Block, BlockState>,
}

impl FrontendBuilderContext {
    /// Creates an empty reusable frontend-builder context.
    pub fn new() -> Self {
        Self::default()
    }

    /// Clears all stored frontend state.
    pub fn clear(&mut self) {
        self.vars.clear();
        self.blocks.clear();
    }
}

/// A frontend-oriented builder for Slonik IR.
pub struct FrontendBuilder<'a> {
    body: &'a mut Body,
    ctx: &'a mut FrontendBuilderContext,
    cur_block: Option<Block>,
}

impl<'a> FrontendBuilder<'a> {
    /// Creates a new frontend-oriented builder.
    pub fn new(body: &'a mut Body, ctx: &'a mut FrontendBuilderContext) -> Self {
        Self {
            body,
            ctx,
            cur_block: None,
        }
    }

    /// Returns a shared reference to the underlying body.
    pub fn body(&self) -> &Body {
        self.body
    }

    /// Returns a mutable reference to the underlying body.
    pub fn body_mut(&mut self) -> &mut Body {
        self.body
    }

    /// Creates a new block.
    pub fn create_block(&mut self) -> Block {
        let block = self.body.create_block();
        self.ctx.blocks.entry(block).or_default();
        block
    }

    /// Switches insertion to `block`.
    pub fn switch_to_block(&mut self, block: Block) {
        self.ctx.blocks.entry(block).or_default();
        self.cur_block = Some(block);
    }

    /// Returns the current insertion block, if any.
    pub fn current_block(&self) -> Option<Block> {
        self.cur_block
    }

    /// Declares a frontend variable with semantic type `ty`.
    pub fn declare_var(&mut self, var: Variable, ty: Type) {
        self.ctx.vars.insert(var, ty);
    }

    /// Returns whether `var` has been declared.
    pub fn is_var_declared(&self, var: Variable) -> bool {
        self.ctx.vars.contains_key(&var)
    }

    /// Returns the declared type of `var`.
    ///
    /// # Panics
    ///
    /// Panics if `var` was not declared first.
    pub fn var_type(&self, var: Variable) -> Type {
        *self
            .ctx
            .vars
            .get(&var)
            .expect("attempted to query undeclared variable")
    }

    /// Returns whether `block` has been sealed.
    pub fn is_sealed(&self, block: Block) -> bool {
        self.ctx
            .blocks
            .get(&block)
            .map(|s| s.sealed)
            .unwrap_or(false)
    }

    /// Returns whether `block` already has a terminator.
    pub fn is_filled(&self, block: Block) -> bool {
        self.ctx
            .blocks
            .get(&block)
            .map(|s| s.filled)
            .unwrap_or(false)
    }

    /// Marks `block` as sealed.
    pub fn seal_block(&mut self, block: Block) {
        self.ctx.blocks.entry(block).or_default().sealed = true;
    }

    /// Marks every known block as sealed.
    pub fn seal_all_blocks(&mut self) {
        for state in self.ctx.blocks.values_mut() {
            state.sealed = true;
        }
    }

    /// Assigns SSA value `value` to frontend variable `var` in the current block.
    pub fn def_var(&mut self, var: Variable, value: Value) {
        let block = self
            .cur_block
            .expect("attempted to define a variable without a current block");

        let expected = self.var_type(var);
        let actual = self.body.value_type(value);

        assert!(
            expected == actual,
            "type mismatch in def_var: variable declared as {expected}, got value of type {actual}",
        );

        self.ctx
            .blocks
            .entry(block)
            .or_default()
            .defs
            .insert(var, value);
    }

    /// Reads the current SSA value for frontend variable `var`.
    pub fn use_var(&mut self, var: Variable) -> Value {
        let block = self
            .cur_block
            .expect("attempted to use a variable without a current block");
        self.use_var_in_block(block, var)
    }

    fn use_var_in_block(&mut self, block: Block, var: Variable) -> Value {
        if let Some(value) = self
            .ctx
            .blocks
            .get(&block)
            .and_then(|state| state.defs.get(&var).copied())
        {
            return value;
        }

        let ty = self.var_type(var);

        let preds = self
            .ctx
            .blocks
            .get(&block)
            .map(|state| state.preds.clone())
            .unwrap_or_default();

        if self.is_sealed(block) && preds.is_empty() {
            panic!(
                "variable {var} used in sealed block {block} with no local definition and no predecessors"
            );
        }

        self.ensure_block_param(block, var, ty)
    }

    fn ensure_block_param(&mut self, block: Block, var: Variable, ty: Type) -> Value {
        if let Some(value) = self
            .ctx
            .blocks
            .get(&block)
            .and_then(|state| state.params.get(&var).copied())
        {
            return value;
        }

        let value = self.body.append_block_param(block, ty);

        {
            let state = self.ctx.blocks.entry(block).or_default();
            state.param_order.push(var);
            state.params.insert(var, value);
            state.defs.insert(var, value);
        }

        let preds = self
            .ctx
            .blocks
            .get(&block)
            .map(|state| state.preds.clone())
            .unwrap_or_default();

        for edge in preds {
            let arg = self.use_var_in_block(edge.pred, var);
            self.append_edge_arg(edge, arg);
        }

        value
    }

    fn append_edge_arg(&mut self, edge: IncomingEdge, arg: Value) {
        let body = self.body_mut();

        let Body {
            stmts, value_lists, ..
        } = body;

        match (&mut stmts[edge.stmt], edge.kind) {
            (StmtData::Jump { dst }, IncomingEdgeKind::Jump) => {
                dst.args.push(arg, value_lists);
            }

            (StmtData::BrIf { then_dst, .. }, IncomingEdgeKind::Then) => {
                then_dst.args.push(arg, value_lists);
            }

            (StmtData::BrIf { else_dst, .. }, IncomingEdgeKind::Else) => {
                else_dst.args.push(arg, value_lists);
            }

            _ => {
                panic!("frontend predecessor bookkeeping became inconsistent");
            }
        }
    }

    fn record_incoming_edge(
        &mut self,
        dst: Block,
        pred: Block,
        stmt: Stmt,
        kind: IncomingEdgeKind,
    ) {
        self.ctx
            .blocks
            .entry(dst)
            .or_default()
            .preds
            .push(IncomingEdge { pred, stmt, kind });
    }

    fn edge_args_to(&mut self, dst: Block) -> Vec<Value> {
        let vars = self
            .ctx
            .blocks
            .get(&dst)
            .map(|state| state.param_order.clone())
            .unwrap_or_default();

        vars.into_iter().map(|var| self.use_var(var)).collect()
    }

    // ------------------------------------------------------------------ //
    // Instruction emission helpers                                         //
    // ------------------------------------------------------------------ //

    fn append_value_inst_in_current_block(
        &mut self,
        ty: Type,
        make: impl FnOnce(Value) -> StmtData,
    ) -> Value {
        let block = self
            .cur_block
            .expect("attempted to build an instruction without a current block");

        assert!(
            !self.is_filled(block),
            "attempted to append an instruction after the block was already terminated"
        );

        self.body.append_value_inst(block, ty, make)
    }

    fn append_stmt_in_current_block(&mut self, data: StmtData) -> Stmt {
        let block = self
            .cur_block
            .expect("attempted to append a statement without a current block");

        assert!(
            !self.is_filled(block),
            "attempted to append a statement after the block was already terminated"
        );

        let is_terminator = data.is_terminator();
        let stmt = self.body.append_stmt(block, data);

        if is_terminator {
            self.ctx.blocks.entry(block).or_default().filled = true;
        }

        stmt
    }

    // ------------------------------------------------------------------ //
    // Value-producing instructions                                         //
    // ------------------------------------------------------------------ //

    /// Appends an integer constant instruction.
    pub fn iconst(&mut self, ty: Type, imm: i64) -> Value {
        self.append_value_inst_in_current_block(ty, |result| StmtData::IConst { imm, result })
    }

    /// Appends a 32-bit float constant instruction from raw IEEE-754 bits.
    pub fn f32const_bits(&mut self, bits: u32) -> Value {
        self.append_value_inst_in_current_block(Type::f32(), |result| StmtData::F32Const {
            bits,
            result,
        })
    }

    /// Appends a 64-bit float constant instruction from raw IEEE-754 bits.
    pub fn f64const_bits(&mut self, bits: u64) -> Value {
        self.append_value_inst_in_current_block(Type::f64(), |result| StmtData::F64Const {
            bits,
            result,
        })
    }

    /// Appends a boolean constant instruction.
    pub fn bconst(&mut self, value: bool) -> Value {
        self.append_value_inst_in_current_block(Type::bool(), |result| StmtData::BConst {
            value,
            result,
        })
    }

    fn unary(&mut self, op: UnaryOp, arg: Value, ty: Type) -> Value {
        self.append_value_inst_in_current_block(ty, |result| StmtData::Unary { op, arg, result })
    }

    fn binary(&mut self, op: BinaryOp, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.append_value_inst_in_current_block(ty, |result| StmtData::Binary {
            op,
            lhs,
            rhs,
            result,
        })
    }

    /// Appends a cast instruction.
    pub fn cast(&mut self, op: CastOp, arg: Value, ty: Type) -> Value {
        self.append_value_inst_in_current_block(ty, |result| StmtData::Cast { op, arg, result })
    }

    /// Appends a bitwise not.  Result type is inferred from the operand.
    pub fn not(&mut self, arg: Value) -> Value {
        let ty = self.body.value_type(arg);
        self.unary(UnaryOp::Not, arg, ty)
    }

    /// Appends an integer comparison.
    pub fn icmp(&mut self, cc: IntCC, lhs: Value, rhs: Value) -> Value {
        self.append_value_inst_in_current_block(Type::bool(), |result| StmtData::Icmp {
            cc,
            lhs,
            rhs,
            result,
        })
    }

    /// Appends a floating-point comparison.
    pub fn fcmp(&mut self, cc: FloatCC, lhs: Value, rhs: Value) -> Value {
        self.append_value_inst_in_current_block(Type::bool(), |result| StmtData::Fcmp {
            cc,
            lhs,
            rhs,
            result,
        })
    }

    /// Appends a select instruction.
    pub fn select(&mut self, cond: Value, if_true: Value, if_false: Value, ty: Type) -> Value {
        self.append_value_inst_in_current_block(ty, |result| StmtData::Select {
            cond,
            if_true,
            if_false,
            result,
        })
    }

    /// Appends a memory load instruction.
    pub fn load(&mut self, addr: Value, size: MemSize, ty: Type) -> Value {
        self.append_value_inst_in_current_block(ty, |result| StmtData::Load { addr, size, result })
    }

    /// Appends an integer add instruction.
    pub fn iadd(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::IAdd, lhs, rhs, ty)
    }

    /// Appends an integer subtract instruction.
    pub fn isub(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::ISub, lhs, rhs, ty)
    }

    /// Appends an integer multiply instruction.
    pub fn imul(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::IMul, lhs, rhs, ty)
    }

    /// Appends a bitwise and instruction.
    pub fn and(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::And, lhs, rhs, ty)
    }

    /// Appends a bitwise or instruction.
    pub fn or(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::Or, lhs, rhs, ty)
    }

    /// Appends a bitwise xor instruction.
    pub fn xor(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::Xor, lhs, rhs, ty)
    }

    /// Appends a logical shift-left instruction.
    pub fn shl(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::Shl, lhs, rhs, ty)
    }

    /// Appends a logical shift-right instruction.
    pub fn lshr(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::LShr, lhs, rhs, ty)
    }

    /// Appends an arithmetic shift-right instruction.
    pub fn ashr(&mut self, lhs: Value, rhs: Value, ty: Type) -> Value {
        self.binary(BinaryOp::AShr, lhs, rhs, ty)
    }

    /// Appends a zero-extension instruction.
    pub fn zext(&mut self, arg: Value, ty: Type) -> Value {
        self.cast(CastOp::Zext, arg, ty)
    }

    /// Appends a sign-extension instruction.
    pub fn sext(&mut self, arg: Value, ty: Type) -> Value {
        self.cast(CastOp::Sext, arg, ty)
    }

    /// Appends a truncation instruction.
    pub fn trunc(&mut self, arg: Value, ty: Type) -> Value {
        self.cast(CastOp::Trunc, arg, ty)
    }

    /// Appends a bitcast instruction.
    pub fn bitcast(&mut self, arg: Value, ty: Type) -> Value {
        self.cast(CastOp::Bitcast, arg, ty)
    }

    // ------------------------------------------------------------------ //
    // Side-effecting and terminator statements                             //
    // ------------------------------------------------------------------ //

    /// Appends a memory store statement.
    pub fn store(&mut self, addr: Value, value: Value, size: MemSize) -> Stmt {
        self.append_stmt_in_current_block(StmtData::Store { addr, value, size })
    }

    /// Appends an effectful call statement.
    pub fn call(&mut self, callee: Value, args: &[Value]) -> Stmt {
        let args = self.body.make_value_list(args);
        self.append_stmt_in_current_block(StmtData::Call { callee, args })
    }

    /// Appends an unconditional jump to `dst`.
    pub fn jump(&mut self, dst: Block) -> Stmt {
        let pred = self
            .cur_block
            .expect("attempted to append jump without a current block");

        let args = self.edge_args_to(dst);
        let dst_call = self.body.block_call(dst, &args);
        let stmt = self.append_stmt_in_current_block(StmtData::Jump { dst: dst_call });

        self.record_incoming_edge(dst, pred, stmt, IncomingEdgeKind::Jump);
        stmt
    }

    /// Appends a conditional branch.
    pub fn br_if(&mut self, cond: Value, then_block: Block, else_block: Block) -> Stmt {
        let pred = self
            .cur_block
            .expect("attempted to append br_if without a current block");

        let then_args = self.edge_args_to(then_block);
        let else_args = self.edge_args_to(else_block);

        let then_dst = self.body.block_call(then_block, &then_args);
        let else_dst = self.body.block_call(else_block, &else_args);

        let stmt = self.append_stmt_in_current_block(StmtData::BrIf {
            cond,
            then_dst,
            else_dst,
        });

        self.record_incoming_edge(then_block, pred, stmt, IncomingEdgeKind::Then);
        self.record_incoming_edge(else_block, pred, stmt, IncomingEdgeKind::Else);

        stmt
    }

    /// Appends a return statement.
    pub fn ret(&mut self, values: &[Value]) -> Stmt {
        let values = self.body.make_value_list(values);
        self.append_stmt_in_current_block(StmtData::Return { values })
    }
}
