use crate::variable_map::{Variable, VariableMap};
use pretty::{DocAllocator, DocBuilder, Pretty};

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

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

        for stmt in self.statements.iter() {
            result = result
                .append(stmt.pretty(variable_map, allocator))
                .append(allocator.text(";"))
                .append(allocator.hardline());
        }

        result
    }
}

#[derive(Debug, PartialEq)]
pub enum Statement<Annotation> {
    Binding(Annotation, Variable, Expression<Annotation>),
    Print(Annotation, Variable),
}

impl<Annotation> Statement<Annotation> {
    pub fn pretty<'a, D, A>(
        &self,
        variable_map: &VariableMap,
        allocator: &'a D,
    ) -> DocBuilder<'a, D, A>
    where
        A: 'a,
        D: ?Sized + DocAllocator<'a, A>,
    {
        match self {
            Statement::Binding(_, var, expr) => {
                let name = variable_map.get_name(*var).unwrap_or("<unknown>");

                allocator
                    .text(name.to_string())
                    .append(allocator.space())
                    .append(allocator.text("="))
                    .append(allocator.space())
                    .append(expr.pretty(variable_map, allocator))
            }
            Statement::Print(_, var) => {
                let name = variable_map.get_name(*var).unwrap_or("<unknown>");

                allocator
                    .text("print")
                    .append(allocator.space())
                    .append(allocator.text(name.to_string()))
            }
        }
    }
}

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

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Primitive {
    Plus,
    Minus,
    Times,
    Divide,
}

impl<Annotation> Expression<Annotation> {
    fn pretty<'a, A, D>(&self, variable_map: &VariableMap, allocator: &'a D) -> DocBuilder<'a, D, A>
    where
        A: 'a,
        D: ?Sized + DocAllocator<'a, A>,
    {
        match self {
            Expression::Value(_, val) => val.pretty(allocator),
            Expression::Reference(_, var) => {
                let name = variable_map.get_name(*var).unwrap_or("<unknown>");

                allocator.text(name.to_string())
            }
            Expression::Primitive(_, op, exprs) if exprs.len() == 1 => op
                .pretty(allocator)
                .append(exprs[0].pretty(variable_map, allocator)),
            Expression::Primitive(_, op, exprs) if exprs.len() == 2 => {
                let left = exprs[0].pretty(variable_map, allocator);
                let right = exprs[1].pretty(variable_map, allocator);

                left.append(allocator.space())
                    .append(op.pretty(allocator))
                    .append(allocator.space())
                    .append(right)
                    .parens()
            }
            Expression::Primitive(_, op, exprs) => {
                allocator.text(format!("!!{:?} with {} arguments!!", op, exprs.len()))
            }
        }
    }
}

impl<'a, 'b, D, A> Pretty<'a, D, A> for &'b Primitive
where
    A: 'a,
    D: ?Sized + DocAllocator<'a, A>,
{
    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
        match self {
            Primitive::Plus => allocator.text("+"),
            Primitive::Minus => allocator.text("-"),
            Primitive::Times => allocator.text("*"),
            Primitive::Divide => allocator.text("/"),
        }
    }
}

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