use crate::syntax::ast::{Expression, Program, Statement, Value};
use pretty::{DocAllocator, DocBuilder, Pretty};

use super::{ConstantType, TopLevel};

impl<'a, 'b, D, A> Pretty<'a, D, A> for &'b Program
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        let mut result = allocator.nil();

        for tl in self.items.iter() {
            result = result
                .append(tl.pretty(allocator))
                .append(allocator.text(";"))
                .append(allocator.hardline());
        }

        result
    }
}

impl<'a, 'b, D, A> Pretty<'a, D, A> for &'b TopLevel
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        match self {
            TopLevel::Statement(stmt) => stmt.pretty(allocator),
            TopLevel::Function(name, arg_names, body) => allocator
                .text("function")
                .append(allocator.space())
                .append(
                    name.as_ref()
                        .map(|x| allocator.text(x.to_string()))
                        .unwrap_or_else(|| allocator.nil()),
                )
                .append(
                    allocator
                        .intersperse(
                            arg_names.iter().map(|x| allocator.text(x.to_string())),
                            CommaSep {},
                        )
                        .parens(),
                )
                .append(allocator.space())
                .append(body.pretty(allocator)),
        }
    }
}

impl<'a, 'b, D, A> Pretty<'a, D, A> for &'b Statement
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        match self {
            Statement::Binding(_, var, expr) => allocator
                .text(var.to_string())
                .append(allocator.space())
                .append(allocator.text("="))
                .append(allocator.space())
                .append(expr.pretty(allocator)),
            Statement::Print(_, var) => allocator
                .text("print")
                .append(allocator.space())
                .append(allocator.text(var.to_string())),
        }
    }
}

impl<'a, 'b, D, A> Pretty<'a, D, A> for &'b Expression
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        match self {
            Expression::Value(_, val) => val.pretty(allocator),
            Expression::Reference(_, var) => allocator.text(var.to_string()),
            Expression::Cast(_, t, e) => allocator
                .text(t.clone())
                .angles()
                .append(e.pretty(allocator)),
            Expression::Primitive(_, op, exprs) if exprs.len() == 1 => allocator
                .text(op.to_string())
                .append(exprs[0].pretty(allocator)),
            Expression::Primitive(_, op, exprs) if exprs.len() == 2 => {
                let left = exprs[0].pretty(allocator);
                let right = exprs[1].pretty(allocator);

                left.append(allocator.space())
                    .append(allocator.text(op.to_string()))
                    .append(allocator.space())
                    .append(right)
                    .parens()
            }
            Expression::Primitive(_, op, exprs) => {
                let call = allocator.text(op.to_string());
                let args = exprs.iter().map(|x| x.pretty(allocator));
                let comma_sepped_args = allocator.intersperse(args, CommaSep {});
                call.append(comma_sepped_args.parens())
            }
            Expression::Block(_, stmts) => match stmts.split_last() {
                None => allocator.text("()"),
                Some((last, &[])) => last.pretty(allocator),
                Some((last, start)) => {
                    let mut result = allocator.text("{").append(allocator.hardline());

                    for stmt in start.iter() {
                        result = result
                            .append(stmt.pretty(allocator))
                            .append(allocator.text(";"))
                            .append(allocator.hardline());
                    }

                    result
                        .append(last.pretty(allocator))
                        .append(allocator.hardline())
                        .append(allocator.text("}"))
                }
            },
        }
    }
}

impl<'a, 'b, D, A> Pretty<'a, D, A> for &'b Value
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        match self {
            Value::Number(opt_base, ty, value) => {
                let value_str = match opt_base {
                    None => format!("{}{}", value, type_suffix(ty)),
                    Some(2) => format!("0b{:b}{}", value, type_suffix(ty)),
                    Some(8) => format!("0o{:o}{}", value, type_suffix(ty)),
                    Some(10) => format!("0d{}{}", value, type_suffix(ty)),
                    Some(16) => format!("0x{:x}{}", value, type_suffix(ty)),
                    Some(_) => format!("!!{:x}{}!!", value, type_suffix(ty)),
                };

                allocator.text(value_str)
            }
        }
    }
}

fn type_suffix(x: &Option<ConstantType>) -> &'static str {
    match x {
        None => "",
        Some(ConstantType::Void) => "<void>",
        Some(ConstantType::I8) => "i8",
        Some(ConstantType::I16) => "i16",
        Some(ConstantType::I32) => "i32",
        Some(ConstantType::I64) => "i64",
        Some(ConstantType::U8) => "u8",
        Some(ConstantType::U16) => "u16",
        Some(ConstantType::U32) => "u32",
        Some(ConstantType::U64) => "u64",
    }
}

#[derive(Clone, Copy)]
struct CommaSep {}

impl<'a, D, A> Pretty<'a, D, A> for CommaSep
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        allocator.text(",").append(allocator.space())
    }
}