use crate::{
    ast::{Constant, Program, Term, Type},
    flat::Binder,
};
use pallas_primitives::babbage::{Constr, PlutusData};
use pretty::RcDoc;
use std::ascii::escape_default;

impl<'a, T> Program<T>
where
    T: Binder<'a>,
{
    pub fn to_pretty(&self) -> String {
        let mut w = Vec::new();

        self.to_doc().render(80, &mut w).unwrap();

        String::from_utf8(w)
            .unwrap()
            .lines()
            // This is a hack to deal with blank newlines
            // that end up with a bunch of useless whitespace
            // because of the nesting
            .map(|l| {
                if l.chars().all(|c| c.is_whitespace()) {
                    "".to_string()
                } else {
                    l.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n")
    }

    fn to_doc(&self) -> RcDoc<()> {
        let version = format!("{}.{}.{}", self.version.0, self.version.1, self.version.2);

        RcDoc::text("(")
            .append(RcDoc::text("program"))
            .append(RcDoc::line())
            .append(RcDoc::text(version))
            .append(RcDoc::line())
            .append(self.term.to_doc())
            .nest(2)
            .append(RcDoc::line_())
            .append(RcDoc::text(")"))
    }
}

impl<'a, T> Term<T>
where
    T: Binder<'a>,
{
    pub fn to_pretty(&self) -> String {
        let mut w = Vec::new();

        self.to_doc().render(80, &mut w).unwrap();

        String::from_utf8(w)
            .unwrap()
            .lines()
            // This is a hack to deal with blank newlines
            // that end up with a bunch of useless whitespace
            // because of the nesting
            .map(|l| {
                if l.chars().all(|c| c.is_whitespace()) {
                    "".to_string()
                } else {
                    l.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n")
    }

    fn to_doc(&self) -> RcDoc<()> {
        match self {
            Term::Var(name) => RcDoc::text(name.text()),
            Term::Delay(term) => RcDoc::text("(")
                .append(
                    RcDoc::text("delay")
                        .append(RcDoc::line())
                        .append(term.to_doc())
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::text(")")),
            Term::Lambda {
                parameter_name,
                body,
            } => RcDoc::text("(")
                .append(
                    RcDoc::text("lam")
                        .append(RcDoc::line())
                        .append(RcDoc::text(parameter_name.text()))
                        .append(RcDoc::line())
                        .append(body.to_doc())
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::text(")")),
            Term::Apply { function, argument } => RcDoc::text("[")
                .append(
                    RcDoc::line()
                        .append(
                            function
                                .to_doc()
                                .append(RcDoc::line())
                                .append(argument.to_doc())
                                .group(),
                        )
                        .nest(2),
                )
                .append(RcDoc::line())
                .append(RcDoc::text("]")),
            Term::Constant(constant) => RcDoc::text("(")
                .append(
                    RcDoc::text("con")
                        .append(RcDoc::line())
                        .append(constant.to_doc())
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::text(")")),
            Term::Force(term) => RcDoc::text("(")
                .append(
                    RcDoc::text("force")
                        .append(RcDoc::line())
                        .append(term.to_doc())
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::text(")")),
            Term::Error => RcDoc::text("(")
                .append(RcDoc::text("error").nest(2))
                .append(RcDoc::line())
                .append(RcDoc::line_())
                .append(RcDoc::text(")")),
            Term::Builtin(builtin) => RcDoc::text("(")
                .append(
                    RcDoc::text("builtin")
                        .append(RcDoc::line())
                        .append(RcDoc::text(builtin.to_string()))
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::text(")")),
            Term::Constr { tag, fields } => RcDoc::text("(")
                .append(
                    RcDoc::text("constr")
                        .append(RcDoc::line())
                        .append(RcDoc::as_string(tag))
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::intersperse(
                    fields.iter().map(|f| f.to_doc()),
                    RcDoc::line_(),
                ))
                .append(RcDoc::text(")")),
            Term::Case { constr, branches } => RcDoc::text("(")
                .append(
                    RcDoc::text("case")
                        .append(RcDoc::line())
                        .append(constr.to_doc())
                        .nest(2),
                )
                .append(RcDoc::line_())
                .append(RcDoc::intersperse(
                    branches.iter().map(|f| f.to_doc()),
                    RcDoc::line_(),
                ))
                .append(RcDoc::text(")")),
        }
        .group()
    }
}

impl Constant {
    pub fn to_pretty(&self) -> String {
        let mut w = Vec::new();

        self.to_doc().render(80, &mut w).unwrap();

        String::from_utf8(w)
            .unwrap()
            .lines()
            // This is a hack to deal with blank newlines
            // that end up with a bunch of useless whitespace
            // because of the nesting
            .map(|l| {
                if l.chars().all(|c| c.is_whitespace()) {
                    "".to_string()
                } else {
                    l.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n")
    }

    fn to_doc(&self) -> RcDoc<()> {
        match self {
            Constant::Integer(i) => RcDoc::text("integer")
                .append(RcDoc::line())
                .append(RcDoc::as_string(i)),
            Constant::ByteString(bs) => RcDoc::text("bytestring")
                .append(RcDoc::line())
                .append(RcDoc::text("#"))
                .append(RcDoc::text(hex::encode(bs))),
            Constant::String(s) => RcDoc::text("string")
                .append(RcDoc::line())
                .append(RcDoc::text("\""))
                .append(RcDoc::text(
                    String::from_utf8(
                        s.as_bytes()
                            .iter()
                            .flat_map(|c| escape_default(*c).collect::<Vec<u8>>())
                            .collect(),
                    )
                    .unwrap(),
                ))
                .append(RcDoc::text("\"")),
            Constant::Unit => RcDoc::text("unit")
                .append(RcDoc::line())
                .append(RcDoc::text("()")),
            Constant::Bool(b) => RcDoc::text("bool")
                .append(RcDoc::line())
                .append(RcDoc::text(if *b { "True" } else { "False" })),
            Constant::ProtoList(r#type, items) => RcDoc::text("(")
                .append("list")
                .append(RcDoc::space())
                .append(r#type.to_doc())
                .append(")")
                .append(RcDoc::line())
                .append(RcDoc::text("["))
                .append(RcDoc::intersperse(
                    items.iter().map(|c| c.to_doc_list()),
                    RcDoc::text(", "),
                ))
                .append(RcDoc::text("]")),
            Constant::ProtoPair(type_left, type_right, left, right) => RcDoc::text("(")
                .append("pair")
                .append(RcDoc::space())
                .append(type_left.to_doc())
                .append(RcDoc::space())
                .append(type_right.to_doc())
                .append(RcDoc::text(")"))
                .append(RcDoc::line())
                .append(RcDoc::text("("))
                .append(left.to_doc_list())
                .append(RcDoc::text(", "))
                .append(right.to_doc_list())
                .append(RcDoc::text(")")),
            Constant::Data(d) => RcDoc::text("data ")
                .append(RcDoc::text("("))
                .append(Self::to_doc_list_plutus_data(d))
                .append(RcDoc::text(")")),
        }
    }

    fn to_doc_list(&self) -> RcDoc<()> {
        match self {
            Constant::Integer(i) => RcDoc::as_string(i),
            Constant::ByteString(bs) => RcDoc::text("#").append(RcDoc::text(hex::encode(bs))),
            Constant::String(s) => RcDoc::text("\"")
                .append(RcDoc::text(s))
                .append(RcDoc::text("\"")),
            Constant::Unit => RcDoc::text("()"),
            Constant::Bool(b) => RcDoc::text(if *b { "True" } else { "False" }),
            Constant::ProtoList(_, items) => RcDoc::text("[")
                .append(RcDoc::intersperse(
                    items.iter().map(|c| c.to_doc_list()),
                    RcDoc::text(", "),
                ))
                .append(RcDoc::text("]")),
            Constant::ProtoPair(_, _, left, right) => RcDoc::text("(")
                .append((*left).to_doc_list())
                .append(RcDoc::text(", "))
                .append((*right).to_doc_list())
                .append(RcDoc::text(")")),

            Constant::Data(data) => Self::to_doc_list_plutus_data(data),
        }
    }

    // This feels a little awkward here; not sure if it should be upstreamed to pallas
    fn to_doc_list_plutus_data(data: &PlutusData) -> RcDoc<()> {
        match data {
            PlutusData::Constr(Constr { tag, fields, .. }) => RcDoc::text("Constr")
                .append(RcDoc::space())
                .append(RcDoc::as_string(tag))
                .append(RcDoc::space())
                .append(RcDoc::text("["))
                .append(RcDoc::intersperse(
                    fields.iter().map(Self::to_doc_list_plutus_data),
                    RcDoc::text(", "),
                ))
                .append(RcDoc::text("]")),
            PlutusData::Map(kvp) => RcDoc::text("Map")
                .append(RcDoc::space())
                .append(RcDoc::text("["))
                .append(RcDoc::intersperse(
                    kvp.iter().map(|(key, value)| {
                        RcDoc::text("(")
                            .append(Self::to_doc_list_plutus_data(key))
                            .append(RcDoc::text(", "))
                            .append(Self::to_doc_list_plutus_data(value))
                            .append(RcDoc::text(")"))
                    }),
                    RcDoc::text(", "),
                ))
                .append(RcDoc::text("]")),
            PlutusData::BigInt(bi) => RcDoc::text("I").append(RcDoc::space()).append(match bi {
                pallas_primitives::babbage::BigInt::Int(v) => RcDoc::text(v.to_string()),
                pallas_primitives::babbage::BigInt::BigUInt(v) => RcDoc::text(v.to_string()),
                pallas_primitives::babbage::BigInt::BigNInt(v) => RcDoc::text(v.to_string()),
            }),
            PlutusData::BoundedBytes(bs) => RcDoc::text("B")
                .append(RcDoc::space())
                .append(RcDoc::text("#"))
                .append(RcDoc::text(hex::encode(bs.to_vec()))),
            PlutusData::Array(a) => RcDoc::text("List")
                .append(RcDoc::space())
                .append(RcDoc::text("["))
                .append(RcDoc::intersperse(
                    a.iter().map(Self::to_doc_list_plutus_data),
                    RcDoc::text(", "),
                ))
                .append(RcDoc::text("]")),
        }
    }
}

impl Type {
    fn to_doc(&self) -> RcDoc<()> {
        match self {
            Type::Bool => RcDoc::text("bool"),
            Type::Integer => RcDoc::text("integer"),
            Type::String => RcDoc::text("string"),
            Type::ByteString => RcDoc::text("bytestring"),
            Type::Unit => RcDoc::text("unit"),
            Type::List(r#type) => RcDoc::text("(list")
                .append(RcDoc::line())
                .append(r#type.to_doc())
                .append(RcDoc::line_())
                .append(")"),
            Type::Pair(l, r) => RcDoc::text("(pair")
                .append(RcDoc::line())
                .append(l.to_doc())
                .append(RcDoc::line())
                .append(r.to_doc())
                .append(")"),
            Type::Data => RcDoc::text("data"),
        }
    }
}

#[cfg(test)]
mod tests {
    use indoc::indoc;
    use pretty_assertions::assert_eq;

    #[test]
    fn format_pair_bool_pair_integer_bytestring() {
        let uplc = "(program 0.0.0 (con (pair bool (pair integer bytestring)) (True, (14, #42))))";

        assert_eq!(
            crate::parser::program(uplc).unwrap().to_pretty(),
            indoc! {
                r#"
                (program
                  0.0.0
                  (con (pair bool (pair integer bytestring)) (True, (14, #42)))
                )"#
            }
        )
    }
}