use crate::syntax::ast::{ConstantType, Expression, Program, TopLevel, Type, Value};
use crate::util::pretty::{derived_display, Allocator};
use pretty::{DocAllocator, DocBuilder};

impl Program {
    pub fn pretty<'a>(&self, allocator: &'a Allocator<'a>) -> DocBuilder<'a, Allocator<'a>> {
        let mut result = allocator.nil();

        for definition in self.structures.values() {
            let mut interior = allocator.nil();

            for (name, ty) in definition.fields.iter() {
                let mut type_bit = allocator.nil();

                if let Some(ty) = ty {
                    type_bit = allocator
                        .text(":")
                        .append(allocator.space())
                        .append(ty.pretty(allocator));
                }

                interior = interior
                    .append(name.original_name().to_string())
                    .append(type_bit)
                    .append(allocator.text(";"))
                    .append(allocator.hardline());
            }

            interior = interior.indent(9);

            let start = allocator
                .text("struct")
                .append(allocator.space())
                .append(allocator.text(definition.name.original_name().to_string()))
                .append(allocator.space())
                .append(allocator.text("{"));

            let conclusion = allocator.text("}").append(allocator.hardline());

            result = result.append(start.append(interior).append(conclusion));
        }

        for definition in self.functions.values() {
            let mut return_type_bit = allocator.nil();

            if let Some(rettype) = definition.return_type.as_ref() {
                return_type_bit = allocator
                    .text("->")
                    .append(allocator.space())
                    .append(rettype.pretty(allocator));
            }

            result = result
                .append(allocator.text("function"))
                .append(allocator.space())
                .append(allocator.text(definition.name.original_name().to_string()))
                .append(allocator.text("("))
                .append(allocator.intersperse(
                    definition.arguments.iter().map(|(x, t)| {
                        let mut type_bit = allocator.nil();

                        if let Some(ty) = t {
                            type_bit = allocator
                                .text(":")
                                .append(allocator.space())
                                .append(ty.pretty(allocator));
                        }

                        allocator
                            .text(x.original_name().to_string())
                            .append(type_bit)
                    }),
                    allocator.text(","),
                ))
                .append(allocator.text(")"))
                .append(return_type_bit)
                .append(allocator.softline())
                .append(definition.body.pretty(allocator))
                .append(allocator.hardline());
        }

        result.append(self.body.pretty(allocator))
    }
}

impl TopLevel {
    pub fn pretty<'a>(&self, allocator: &'a Allocator<'a>) -> DocBuilder<'a, Allocator<'a>> {
        match self {
            TopLevel::Expression(expr) => expr
                .pretty(allocator)
                .append(allocator.text(";"))
                .append(allocator.hardline()),
            TopLevel::Structure(_, name, fields) => allocator
                .text("struct")
                .append(allocator.space())
                .append(allocator.text(name.to_string()))
                .append(allocator.space())
                .append(allocator.text("{"))
                .append(allocator.hardline())
                .append(
                    allocator
                        .concat(fields.iter().map(|(name, ty)| {
                            let type_bit = if let Some(ty) = ty {
                                allocator
                                    .text(":")
                                    .append(allocator.space())
                                    .append(ty.pretty(allocator))
                            } else {
                                allocator.nil()
                            };
                            allocator
                                .text(name.to_string())
                                .append(type_bit)
                                .append(allocator.text(";"))
                                .append(allocator.hardline())
                        }))
                        .indent(2),
                )
                .append(allocator.text("}"))
                .append(allocator.hardline()),
        }
    }
}

impl Expression {
    pub fn pretty<'a>(&self, allocator: &'a Allocator<'a>) -> DocBuilder<'a, Allocator<'a>> {
        match self {
            Expression::Value(_, val) => val.pretty(allocator),
            Expression::Constructor(_, name, fields) => allocator
                .text(name.to_string())
                .append(allocator.space())
                .append(allocator.text("{"))
                .append(allocator.hardline())
                .append(
                    allocator
                        .concat(fields.iter().map(|(n, e)| {
                            allocator
                                .text(n.to_string())
                                .append(allocator.text(":"))
                                .append(allocator.space())
                                .append(e.pretty(allocator))
                                .append(allocator.text(";"))
                                .append(allocator.hardline())
                        }))
                        .indent(2),
                )
                .append(allocator.text("}")),
            Expression::Reference(var) => allocator.text(var.to_string()),
            Expression::FieldRef(_, val, field) => val
                .pretty(allocator)
                .append(allocator.text("."))
                .append(allocator.text(field.to_string())),
            Expression::Cast(_, t, e) => allocator
                .text(t.clone())
                .angles()
                .append(e.pretty(allocator)),
            Expression::Primitive(_, op) => allocator.text(op.original_name().to_string()),
            Expression::Call(_, fun, args) => {
                let mut args = args.iter().map(|x| x.pretty(allocator)).collect::<Vec<_>>();

                match fun.as_ref() {
                    Expression::Primitive(_, name)
                        if ["/", "*", "+", "-"].contains(&name.current_name())
                            && args.len() == 2 =>
                    {
                        let second = args.pop().unwrap();
                        args.pop()
                            .unwrap()
                            .append(allocator.space())
                            .append(allocator.text(name.current_name().to_string()))
                            .append(allocator.space())
                            .append(second)
                            .parens()
                    }

                    Expression::Primitive(_, name)
                        if ["negate"].contains(&name.current_name()) && args.len() == 1 =>
                    {
                        allocator.text("-").append(args.pop().unwrap())
                    }

                    Expression::Primitive(_, name)
                        if ["print"].contains(&&name.current_name()) && args.len() == 1 =>
                    {
                        allocator
                            .text("print")
                            .append(allocator.space())
                            .append(args.pop().unwrap())
                    }

                    _ => {
                        let comma_sepped_args = allocator.intersperse(args, allocator.text(","));
                        fun.pretty(allocator).append(comma_sepped_args.parens())
                    }
                }
            }
            Expression::Block(_, stmts) => match stmts.split_last() {
                None => allocator.text("()"),
                Some((last, &[])) => last.pretty(allocator),
                Some((last, start)) => {
                    let beginning = allocator.text("{").append(allocator.hardline());
                    let mut inner = allocator.nil();

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

                    inner = inner
                        .append(last.pretty(allocator))
                        .append(allocator.hardline());

                    inner = inner.indent(2);

                    beginning.append(inner).append(allocator.text("}"))
                }
            },
            Expression::Binding(_, var, expr) => allocator
                .text(var.to_string())
                .append(allocator.space())
                .append(allocator.text("="))
                .append(allocator.space())
                .append(expr.pretty(allocator)),
            Expression::Function(_, name, args, rettype, 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(
                            args.iter().map(|(x, t)| {
                                allocator.text(x.to_string()).append(
                                    t.as_ref()
                                        .map(|t| {
                                            allocator
                                                .text(":")
                                                .append(allocator.space())
                                                .append(t.pretty(allocator))
                                        })
                                        .unwrap_or_else(|| allocator.nil()),
                                )
                            }),
                            allocator.text(","),
                        )
                        .parens(),
                )
                .append(
                    rettype
                        .as_ref()
                        .map(|rettype| {
                            allocator
                                .space()
                                .append(allocator.text("->"))
                                .append(allocator.space())
                                .append(rettype.pretty(allocator))
                        })
                        .unwrap_or_else(|| allocator.nil()),
                )
                .append(allocator.space())
                .append(body.pretty(allocator)),
        }
    }
}

impl Value {
    pub fn pretty<'a>(&self, allocator: &'a Allocator<'a>) -> DocBuilder<'a, Allocator<'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)
            }

            Value::Void => allocator.text("void()"),
        }
    }
}

fn type_suffix(x: &Option<ConstantType>) -> &'static str {
    match x {
        None => "",
        Some(ConstantType::Void) => panic!("Should never get a void type suffix."),
        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",
    }
}

impl Type {
    pub fn pretty<'a>(&self, allocator: &'a Allocator<'a>) -> DocBuilder<'a, Allocator<'a>> {
        match self {
            Type::Named(x) => allocator.text(x.to_string()),
            Type::Struct(fields) => allocator
                .text("struct")
                .append(allocator.space())
                .append(allocator.intersperse(
                    fields.iter().map(|(name, ty)| {
                        allocator
                            .text(name.to_string())
                            .append(allocator.text(":"))
                            .append(allocator.space())
                            .append(
                                ty.as_ref()
                                    .map(|x| x.pretty(allocator))
                                    .unwrap_or_else(|| allocator.text("_")),
                            )
                            .append(allocator.text(";"))
                    }),
                    allocator.hardline(),
                ).braces())
        }
    }
}

derived_display!(Program);
derived_display!(TopLevel);
derived_display!(Expression);
derived_display!(Value);