use crate::syntax::token_stream::Location;
use pretty::{DocAllocator, DocBuilder, Pretty};

static BINARY_OPERATORS: &[&str] = &["+", "-", "*", "/"];

#[derive(Debug, PartialEq)]
pub struct Program {
    pub statements: Vec<Statement>,
}

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 stmt in self.statements.iter() {
            result = result
                .append(stmt.pretty(allocator))
                .append(allocator.text(";"))
                .append(allocator.hardline());
        }

        result
    }
}

#[derive(Debug, PartialEq)]
pub enum Statement {
    Binding(Location, String, Expression),
    Print(Location, String),
    Expr(Location, Expression),
}

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())),
            Statement::Expr(_, expr) => expr.pretty(allocator),
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Expression {
    Value(Location, Value),
    Reference(Location, String),
    Primitive(Location, String, Vec<Expression>),
}

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::Primitive(_, op, exprs) if BINARY_OPERATORS.contains(&op.as_ref()) => {
                assert_eq!(
                    exprs.len(),
                    2,
                    "Found binary operator with {} components?",
                    exprs.len()
                );

                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())
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Value {
    Number(Option<u8>, i128),
}

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, value) => {
                let value_str = match opt_base {
                    None => format!("{}", value),
                    Some(2) => format!("0b{:b}", value),
                    Some(8) => format!("0o{:o}", value),
                    Some(10) => format!("0d{}", value),
                    Some(16) => format!("0x{:x}", value),
                    Some(_) => format!("!!{:x}!!", value),
                };

                allocator.text(value_str)
            }
        }
    }
}

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

impl<'a, 'b, 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())
    }
}