use super::typed_ir;
use super::Context;

impl Context {
    pub fn pretty_typed_block(
        &self,
        w: &mut dyn std::fmt::Write,
        block: &typed_ir::Body,
    ) -> std::fmt::Result {
        let callsig = self.callables.get(block.signature);
        {
            // TODO: ugly copy paste job
            let args = callsig
                .args
                .iter()
                .map(|(_name, typ)| {
                    let mut out = String::new();
                    self.types.pretty_type(&mut out, *typ)?;
                    Ok(out)
                })
                .collect::<Result<Vec<String>, std::fmt::Error>>()?;
            let genargs = callsig
                .genargs
                .iter()
                .map(|(_name, typ)| {
                    let mut type_str = String::new();
                    self.types.pretty_type(&mut type_str, *typ)?;
                    Ok(type_str)
                })
                .collect::<Result<Vec<String>, std::fmt::Error>>()?;
            writeln!(
                w,
                "block {}<{}>({})",
                callsig.name(),
                genargs.join(", "),
                args.join(", ")
            )?;
        }
        for sig in &block.signals {
            let mut typ_pretty = String::new();
            self.types.pretty_type(&mut typ_pretty, sig.typ)?;
            writeln!(w, "sig_{}: {}", sig.id.0, typ_pretty)?
        }
        self.pretty_typed_expr(w, &block.expr)?;
        Ok(())
    }

    pub fn pretty_typed_expr(
        &self,
        w: &mut dyn std::fmt::Write,
        expr: &typed_ir::Expr,
    ) -> std::fmt::Result {
        let expr_pretty = match &expr.kind {
            typed_ir::ExprKind::Literal(lit) => {
                let mut lit_str = String::new();
                self.types.pretty_value(&mut lit_str, lit)?;
                lit_str
            }
            typed_ir::ExprKind::Path(path) => format!("sig_{}", path.0),
            typed_ir::ExprKind::Call(call) => {
                let args = call
                    .args
                    .iter()
                    .map(|arg| {
                        self.pretty_typed_expr(w, arg)?;
                        Ok(format!("_{}", arg.id.0))
                    })
                    .collect::<Result<Vec<_>, std::fmt::Error>>()?;
                let callable = self.callables.get(call.called);
                let genargs = call
                    .genargs
                    .iter()
                    .map(|param| {
                        let mut type_str = String::new();
                        self.types.pretty_type(&mut type_str, *param)?;
                        Ok(type_str)
                    })
                    .collect::<Result<Vec<_>, std::fmt::Error>>()?;
                format!(
                    "{}<{}>({})",
                    callable.name(),
                    genargs.join(", "),
                    args.join(", ")
                )
            }
            typed_ir::ExprKind::Match(match_) => {
                self.pretty_typed_expr(w, &match_.expr)?;
                let arms = match_
                    .arms
                    .iter()
                    .map(|(pat, val)| {
                        self.pretty_typed_expr(w, pat)?;
                        self.pretty_typed_expr(w, val)?;
                        Ok(format!("   _{} =>  _{}", pat.id.0, val.id.0))
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                format!(
                    "match (_{}) {{\n{}\n}}",
                    &match_.expr.id.0,
                    arms.join(",\n")
                )
            }
        };
        let mut type_pretty = String::new();
        self.types.pretty_type(&mut type_pretty, expr.typ)?;
        writeln!(w, "let _{}: {} = {}", expr.id.0, type_pretty, expr_pretty)?;
        Ok(())
    }
}