//! Text formatting for Slonik IR.
//!
//! Expressions that are never referenced by any statement are emitted in a
//! trailing `// dead expressions` section.

use std::collections::HashSet;
use std::fmt;

use crate::ir::{Block, BlockCall, Body, Expr, ExprData, Stmt, StmtData, Value, ValueDef};

/// Writes `body` in textual form.
pub fn write_body(f: &mut fmt::Formatter<'_>, body: &Body) -> fmt::Result {
    let mut printer = Printer::new(body);
    printer.write_body(f)
}

impl fmt::Display for Body {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write_body(f, self)
    }
}

struct Printer<'a> {
    body: &'a Body,
    emitted: HashSet<Value>,
    visiting: HashSet<Expr>,
}

impl<'a> Printer<'a> {
    fn new(body: &'a Body) -> Self {
        Self {
            body,
            emitted: HashSet::new(),
            visiting: HashSet::new(),
        }
    }

    fn write_body(&mut self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "body {{")?;

        let mut first_block = true;
        for block in self.body.blocks() {
            if !first_block {
                writeln!(f)?;
            }
            first_block = false;
            self.write_block(f, block)?;
        }

        let mut wrote_dead_header = false;
        for (expr, node) in self.body.exprs.iter() {
            let value = node.value;
            if !self.emitted.contains(&value) {
                if !wrote_dead_header {
                    if !first_block {
                        writeln!(f)?;
                    }
                    writeln!(f, "  // dead expressions")?;
                    wrote_dead_header = true;
                }

                self.emit_expr(f, expr, 1)?;
            }
        }

        writeln!(f, "}}")
    }

    fn write_block(&mut self, f: &mut fmt::Formatter<'_>, block: Block) -> fmt::Result {
        self.write_block_header(f, block)?;

        for &stmt in self.body.block_stmts(block) {
            self.emit_stmt_deps(f, stmt, 2)?;
            self.write_stmt(f, stmt, 2)?;
        }

        Ok(())
    }

    fn write_block_header(&mut self, f: &mut fmt::Formatter<'_>, block: Block) -> fmt::Result {
        write!(f, "  ^{}", block)?;

        let params = self.body.block_params(block);
        if !params.is_empty() {
            write!(f, "(")?;
            for (i, &param) in params.iter().enumerate() {
                if i != 0 {
                    write!(f, ", ")?;
                }

                let ty = self.body.value_type(param);
                write!(f, "%{}: {}", param, ty)?;
            }
            write!(f, ")")?;
        }

        writeln!(f, ":")
    }

    fn emit_stmt_deps(
        &mut self,
        f: &mut fmt::Formatter<'_>,
        stmt: Stmt,
        indent: usize,
    ) -> fmt::Result {
        match self.body.stmt_data(stmt) {
            StmtData::Store { addr, value, .. } => {
                self.emit_value(f, *addr, indent)?;
                self.emit_value(f, *value, indent)?;
            }

            StmtData::Call { callee, args } => {
                self.emit_value(f, *callee, indent)?;
                for &arg in args.as_slice(&self.body.value_lists) {
                    self.emit_value(f, arg, indent)?;
                }
            }

            StmtData::Jump { dst } => {
                self.emit_block_call_args(f, *dst, indent)?;
            }

            StmtData::BrIf {
                cond,
                then_dst,
                else_dst,
            } => {
                self.emit_value(f, *cond, indent)?;
                self.emit_block_call_args(f, *then_dst, indent)?;
                self.emit_block_call_args(f, *else_dst, indent)?;
            }

            StmtData::Return { values } => {
                for &value in values.as_slice(&self.body.value_lists) {
                    self.emit_value(f, value, indent)?;
                }
            }
        }

        Ok(())
    }

    fn emit_block_call_args(
        &mut self,
        f: &mut fmt::Formatter<'_>,
        call: BlockCall,
        indent: usize,
    ) -> fmt::Result {
        for &arg in call.args.as_slice(&self.body.value_lists) {
            self.emit_value(f, arg, indent)?;
        }
        Ok(())
    }

    fn emit_value(
        &mut self,
        f: &mut fmt::Formatter<'_>,
        value: Value,
        indent: usize,
    ) -> fmt::Result {
        if self.emitted.contains(&value) {
            return Ok(());
        }

        match self.body.value_def(value) {
            ValueDef::Param(_, _) => Ok(()),
            ValueDef::Expr(expr) => self.emit_expr(f, expr, indent),
        }
    }

    fn emit_expr(&mut self, f: &mut fmt::Formatter<'_>, expr: Expr, indent: usize) -> fmt::Result {
        let value = self.body.expr_value(expr);
        if self.emitted.contains(&value) {
            return Ok(());
        }

        if !self.visiting.insert(expr) {
            return Err(fmt::Error);
        }

        match self.body.expr_data(expr) {
            ExprData::IConst { .. }
            | ExprData::F32Const { .. }
            | ExprData::F64Const { .. }
            | ExprData::BConst { .. } => {}

            ExprData::Unary { arg, .. } | ExprData::Cast { arg, .. } => {
                self.emit_value(f, *arg, indent)?;
            }

            ExprData::Binary { lhs, rhs, .. }
            | ExprData::Icmp { lhs, rhs, .. }
            | ExprData::Fcmp { lhs, rhs, .. } => {
                self.emit_value(f, *lhs, indent)?;
                self.emit_value(f, *rhs, indent)?;
            }

            ExprData::Select {
                cond,
                if_true,
                if_false,
            } => {
                self.emit_value(f, *cond, indent)?;
                self.emit_value(f, *if_true, indent)?;
                self.emit_value(f, *if_false, indent)?;
            }

            ExprData::Load { addr, .. } => {
                self.emit_value(f, *addr, indent)?;
            }
        }

        self.write_indent(f, indent)?;
        self.write_expr_binding(f, expr)?;
        self.visiting.remove(&expr);
        self.emitted.insert(value);

        Ok(())
    }

    fn write_expr_binding(&mut self, f: &mut fmt::Formatter<'_>, expr: Expr) -> fmt::Result {
        let value = self.body.expr_value(expr);
        let ty = self.body.value_type(value);

        write!(f, "%{} = ", value)?;

        match self.body.expr_data(expr) {
            ExprData::IConst { imm } => {
                write!(f, "iconst {}", imm)?;
                writeln!(f, " : {}", ty)
            }

            ExprData::F32Const { bits } => {
                write!(f, "f32const 0x{bits:08x}")?;
                writeln!(f, " : {}", ty)
            }

            ExprData::F64Const { bits } => {
                write!(f, "f64const 0x{bits:016x}")?;
                writeln!(f, " : {}", ty)
            }

            ExprData::BConst { value } => {
                write!(f, "bconst {}", value)?;
                writeln!(f, " : {}", ty)
            }

            ExprData::Unary { op, arg } => {
                write!(f, "{} %{}", op, arg)?;
                writeln!(f, " : {}", ty)
            }

            ExprData::Binary { op, lhs, rhs } => {
                write!(f, "{} %{}, %{}", op, lhs, rhs)?;
                writeln!(f, " : {}", ty)
            }

            ExprData::Cast { op, arg } => {
                let src_ty = self.body.value_type(*arg);
                writeln!(f, "{} %{} : {} -> {}", op, arg, src_ty, ty)
            }

            ExprData::Icmp { cc, lhs, rhs } => {
                let operand_ty = self.body.value_type(*lhs);
                writeln!(f, "icmp {} %{}, %{} : {}", cc, lhs, rhs, operand_ty)
            }

            ExprData::Fcmp { cc, lhs, rhs } => {
                let operand_ty = self.body.value_type(*lhs);
                writeln!(f, "fcmp {} %{}, %{} : {}", cc, lhs, rhs, operand_ty)
            }

            ExprData::Select {
                cond,
                if_true,
                if_false,
            } => {
                write!(f, "select %{}, %{}, %{}", cond, if_true, if_false)?;
                writeln!(f, " : {}", ty)
            }

            ExprData::Load { addr, size } => {
                writeln!(f, "load %{} : {} -> {}", addr, size, ty)
            }
        }
    }

    fn write_stmt(&mut self, f: &mut fmt::Formatter<'_>, stmt: Stmt, indent: usize) -> fmt::Result {
        self.write_indent(f, indent)?;

        match self.body.stmt_data(stmt) {
            StmtData::Store { addr, value, size } => {
                writeln!(f, "store %{}, %{} : {}", value, addr, size)
            }

            StmtData::Call { callee, args } => {
                write!(f, "call %{}", callee)?;

                if !args.is_empty() {
                    write!(f, "(")?;
                    self.write_value_list(f, args.as_slice(&self.body.value_lists))?;
                    writeln!(f, ")")?;
                }
                Ok(())
            }

            StmtData::Jump { dst } => {
                write!(f, "br ")?;
                self.write_block_call(f, *dst)?;
                writeln!(f)
            }

            StmtData::BrIf {
                cond,
                then_dst,
                else_dst,
            } => {
                write!(f, "cond_br %{}, ", cond)?;
                self.write_block_call(f, *then_dst)?;
                write!(f, ", ")?;
                self.write_block_call(f, *else_dst)?;
                writeln!(f)
            }

            StmtData::Return { values } => {
                write!(f, "return")?;
                let values = values.as_slice(&self.body.value_lists);
                if !values.is_empty() {
                    write!(f, " ")?;
                    self.write_value_list(f, values)?;
                }
                writeln!(f)
            }
        }
    }

    fn write_block_call(&mut self, f: &mut fmt::Formatter<'_>, call: BlockCall) -> fmt::Result {
        write!(f, "^{}", call.block)?;
        if !call.args.is_empty() {
            write!(f, "(")?;
            self.write_value_list(f, call.args.as_slice(&self.body.value_lists))?;
            write!(f, ")")?;
        }
        Ok(())
    }

    fn write_value_list(&mut self, f: &mut fmt::Formatter<'_>, values: &[Value]) -> fmt::Result {
        for (i, value) in values.iter().enumerate() {
            if i != 0 {
                write!(f, ", ")?;
            }
            write!(f, "%{}", value)?;
        }
        Ok(())
    }

    fn write_indent(&mut self, f: &mut fmt::Formatter<'_>, indent: usize) -> fmt::Result {
        for _ in 0..indent {
            write!(f, "  ")?;
        }
        Ok(())
    }
}