use python_ast::{Module, stmt, expr, boolop, operator, unaryop, cmpop, arguments, arg, keyword, comprehension, withitem, excepthandler, slice, name_constant};

use std::iter;

trait ToStringable {
    fn to_string(&self) -> String;
}

impl boolop {
    fn to_string(&self) -> &'static str {
        match *self {
            boolop::And => " and ",
            boolop::Or => " or "
        }
    }
}

impl operator {
    fn to_string(&self) -> &'static str {
        match *self {
            operator::Add => "+",
            operator::Sub => "-",
            operator::Mult => "*",
            operator::MatMult => "@",
            operator::Div => "/",
            operator::Mod => "%",
            operator::Pow => "**",
            operator::LShift => "<<",
            operator::RShift => ">>",
            operator::BitOr => "|",
            operator::BitXor => "^",
            operator::BitAnd => "&",
            operator::FloorDiv => "//",
        }
    }
}

impl unaryop {
    fn to_string(&self) -> &'static str {
        match *self {
            unaryop::Invert => "~",
            unaryop::Not => "not ",
            unaryop::UAdd => "+",
            unaryop::USub => "-"
        }
    }
}

impl cmpop {
    fn to_string(&self) -> &'static str {
        match *self {
            cmpop::Eq => "==",
            cmpop::NotEq => "!=",
            cmpop::Lt => "<",
            cmpop::LtE => "<=",
            cmpop::Gt => ">",
            cmpop::GtE => ">=",
            cmpop::Is => "is",
            cmpop::IsNot => "is not",
            cmpop::In => "in",
            cmpop::NotIn => "not in",
        }
    }
}

impl ToStringable for withitem {
    fn to_string(&self) -> String {
        match self.optional_vars {
            None => format!("{}", self.context_expr.to_string()),
            Some(ref optional_var) => format!("{} as {}", self.context_expr.to_string(), optional_var.to_string())
        }
    }
}

fn vec_to_strings_vec<T: ToStringable>(values: Vec<T>) -> Vec<String> {
    let mut vector = vec!();
    for value in values {
        vector.push(value.to_string());
    }
    vector
}

impl ToStringable for keyword {
    fn to_string(&self) -> String {
        match self.arg.clone() {
            Some(arg) => format!("{}={}", arg, self.value.to_string()),
            None => format!("**{}", self.value.to_string())
        }
    }
}

fn statements_to_string(indent: usize, body: Vec<stmt>) -> String {
    let mut statements = vec!();
    for statement in body {
        statements.push(statement.to_string(indent + 1));
    }
    statements.join("\n")
}

fn if_else_statements_to_string(indent: usize, body: Vec<stmt>, orelse: Vec<stmt>) -> String {
    let body = statements_to_string(indent, body);
    if orelse.is_empty() {
        body
    } else {
        format!("{}\n{}else:\n{}", body, make_indent(indent), statements_to_string(indent, orelse))
    }
}

impl ToStringable for comprehension {
    fn to_string(&self) -> String {
        let mut result = vec!();
        result.push(format!("for {} in {}", self.target.to_string(), self.iter.to_string()));
        for if_ in self.ifs.clone() {
            result.push(format!("if {}", if_.to_string()))
        }
        result.join(" ")
    }
}

impl ToStringable for slice {
    fn to_string(&self) -> String {
        match *self {
            slice::Index(ref value) => value.to_string(),
            slice::Slice(ref lower, ref upper, ref step) => {
                format!("{}{}{}",
                    match *lower {
                        None => String::new(),
                        Some(ref lower) => lower.to_string()
                    },
                    match *upper {
                        None => String::new(),
                        Some(ref upper) => format!(":{}", upper.to_string())
                    },
                    match *step {
                        None => String::new(),
                        Some(ref step) => format!(":{}", step.to_string())
                    },
                )
            },
            slice::ExtSlice(ref dims) => {
                let mut dims_ = vec!();
                for dim in dims {
                    dims_.push(dim.to_string());
                }
                dims_.join(", ")
            }
        }
    }
}

impl excepthandler {
    fn to_string(self, indent: usize) -> String {
        let current_indent = make_indent(indent);
        format!("{}except {}{}:\n{}",
            current_indent,
            match self.type_ {
                None => String::new(),
                Some(ref type_) => type_.to_string()
            },
            match self.name {
                None => String::new(),
                Some(ref name) => format!(" as {}", name)
            },
            statements_to_string(indent, self.body)
        )
    }
}

impl ToStringable for expr {
    fn to_string(&self) -> String {
        match self.clone() {
            expr::BoolOp(op, values) => {
                let mut data = vec!();
                for value in values {
                    data.push(value.to_string());
                }
                data.join(op.to_string())
            },
            expr::BinOp(a, op, b) => format!("{} {} {}", a.to_string(), op.to_string(), b.to_string()),
            expr::UnaryOp(op, operand) => format!("{}{}", op.to_string(), operand.to_string()),

            expr::Compare(left, ops, comparators) => format!("{} {}", left.to_string(), {
                let mut arguments = vec!();
                for (op, comparator) in ops.iter().zip(comparators.iter()) {
                    arguments.push(format!("{} {}", op.to_string(), comparator.to_string()))
                }
                arguments.join(" ")
            }),
            expr::Call(func, args, keywords) => format!("{}({})", func.to_string(), {
                let mut arguments = vec!();
                let args = vec_to_strings_vec(args);
                let keywords = vec_to_strings_vec(keywords);
                arguments.extend(args);
                arguments.extend(keywords);
                arguments.join(", ")
            }),
            expr::Num(n) => format!("{}", n),
            expr::Str(s) => format!("\"{}\"", s.replace('\\', "\\\\").replace('"', "\\\"")),
            expr::Bytes(s) => format!("b\"{}\"", {
                let mut string = String::with_capacity(s.len());
                for ascii_code in s {
                    let c = ascii_code as char;
                    if c >= ' ' && c <= '~' {
                        if c == '"' || c == '\\' {
                            string.push('\\');
                        }
                        string.push(c);
                    } else if c == '\t' {
                        string.push_str("\\t");
                    } else if c == '\n' {
                        string.push_str("\\n");
                    } else if c == '\r' {
                        string.push_str("\\r");
                    } else /* if c > '~' */ {
                        let value = format!("\\x{:02x}", ascii_code);
                        string.push_str(value.as_str());
                    }
                }
                string
            }),
            expr::NameConstant(constant) => match constant {
                name_constant::True => String::from("True"),
                name_constant::False => String::from("False"),
                name_constant::None_ => String::from("None")
            },
            expr::Attribute(value, attr, _) => format!("{}.{}", value.to_string(), attr),
            expr::Name(name, _) => format!("{}", name),
            expr::List(elts, _) => format!("[{}]", vec_to_strings_vec(elts).join(", ")),
            expr::ListComp(elt, generators) => format!("[{} {}]", elt.to_string(), vec_to_strings_vec(generators).join(" ")),
            expr::DictComp(key, value, generators) => format!("{{{}: {} {}}}", key.to_string(), value.to_string(), vec_to_strings_vec(generators).join(" ")),
            expr::Tuple(elts, _) => format!("({})", vec_to_strings_vec(elts).join(", ")),
            expr::Ellipsis => format!("..."),
            expr::Await(value) => format!("await {}", value.to_string()),
            expr::Yield(value) => {
                match *value {
                    None => String::from("yield"),
                    Some(value) => format!("yield {}", value.to_string())
                }
            },
            expr::YieldFrom(value) => format!("yield from {}", value.to_string()),
            expr::Set(elts) => format!("{{{}}}", vec_to_strings_vec(elts).join(", ")),
            expr::SetComp(elt, generators) => format!("{{{} {}}}", elt.to_string(), vec_to_strings_vec(generators).join(" ")),
            expr::GeneratorExp(elt, generators) => format!("({} {})", elt.to_string(), vec_to_strings_vec(generators).join(" ")),
            expr::Lambda(args, body) => format!("lambda {}: {}", args.to_string(), body.to_string()),
            expr::IfExp(test, body, orelse) => format!("{} if {} else {}", body.to_string(), test.to_string(), orelse.to_string()),
            expr::Dict(keys, values) => {
                format!("{{{}}}",
                    {
                        let mut items = vec!();
                        for (key, value) in keys.iter().zip(values.iter()) {
                            let item = match *key {
                                None => format!("**{}", value.to_string()),
                                Some(ref key) => format!("{}: {}", key.to_string(), value.to_string()),
                            };
                            items.push(item);
                        }
                        items.join(", ")
                    }
                )
            },
            expr::Subscript(value, slice, _) => format!("{}[{}]", value.to_string(), slice.to_string()),
            expr::Starred(value, _) => format!("*{}", value.to_string()),
        }
    }
}

impl ToStringable for arg {
    fn to_string(&self) -> String {
        match self.annotation {
            None => self.arg.clone(),
            Some(ref annotation) => format!("{}: {}", self.arg, annotation.to_string())
        }
    }
}

impl ToStringable for arguments {
    fn to_string(&self) -> String {
        let mut args = vec!();
        let num_args = self.args.len();
        let num_defaults = self.defaults.len();

        let mut normal_args = self.args.clone();
        let defaulted_args = normal_args.split_off(num_args - num_defaults);
        for arg in normal_args {
            args.push(arg.to_string());
        }
        for (arg, default) in defaulted_args.iter().zip(self.defaults.iter()) {
            args.push(format!("{}={}", arg.to_string(), default.to_string()));
        }

        if !self.kwonlyargs.is_empty() {
            match self.vararg {
                None => args.push("*".to_string()),
                Some(ref vararg) => args.push(format!("*{}", vararg.to_string())),
            }
        }
        for (arg, default) in self.kwonlyargs.iter().zip(self.kw_defaults.iter()) {
            let arg = arg.to_string();
            match *default {
                Some(ref default) => args.push(format!("{}={}", arg, default.to_string())),
                None => args.push(arg)
            }
        }
        match self.kwarg {
            Some(ref kwarg) => args.push(format!("**{}", kwarg.to_string())),
            None => ()
        }
        args.join(", ")
    }
}

fn make_indent(indent: usize) -> String {
    iter::repeat("    ").take(indent).collect()
}

fn decorator_list_to_string(decorator_list: Vec<expr>, indent: usize) -> String {
    let current_indent = make_indent(indent);
    let decorators = vec_to_strings_vec(decorator_list);
    if decorators.is_empty() {
        String::new()
    } else {
        format!("{}@{}\n",
            current_indent,
            decorators.join(format!("\n{}@", current_indent).as_str())
        )
    }
}

impl stmt {
    fn to_string(&self, indent: usize) -> String {
        let current_indent = make_indent(indent);
        match self.clone() {
            stmt::ClassDef(name, bases, keywords, body, decorator_list) => {
                format!("{}{}class {}({}):\n{}",
                    decorator_list_to_string(decorator_list, indent),
                    current_indent,
                    name,
                    {
                        let mut arguments = vec!();
                        let bases = vec_to_strings_vec(bases);
                        let keywords = vec_to_strings_vec(keywords);
                        arguments.extend(bases);
                        arguments.extend(keywords);
                        arguments.join(", ")
                    },
                    statements_to_string(indent, body)
                )
            },
            stmt::FunctionDef(name, arguments, body, decorator_list, returns, async) => {
                format!("{}{}{}def {}({}){}:\n{}",
                    decorator_list_to_string(decorator_list, indent),
                    if async { "async " } else { "" },
                    current_indent,
                    name,
                    arguments.to_string(),
                    match returns {
                        None => String::new(),
                        Some(returns) => format!(" -> {}", returns.to_string())
                    },
                    statements_to_string(indent, body)
                )
            }
            stmt::Global(names) => format!("{}global {}", current_indent, names.join(", ")),
            stmt::Nonlocal(names) => format!("{}nonlocal {}", current_indent, names.join(", ")),
            stmt::If(test, body, orelse) => format!("{}if {}:\n{}", current_indent, test.to_string(), if_else_statements_to_string(indent, body, orelse)),
            stmt::While(test, body, orelse) => format!("{}while {}:\n{}", current_indent, test.to_string(), if_else_statements_to_string(indent, body, orelse)),
            stmt::For(target, iter, body, orelse, async) => {
                format!("{}{}for {} in {}:\n{}",
                    current_indent,
                    if async { "async " } else { "" },
                    target.to_string(),
                    iter.to_string(),
                    if_else_statements_to_string(indent, body, orelse)
                )
            },
            stmt::Assign(targets, value) => format!("{}{} = {}", current_indent, vec_to_strings_vec(targets).join(" = "), value.to_string()),
            stmt::AugAssign(target, op, value) => format!("{}{} {}= {}", current_indent, target.to_string(), op.to_string(), value.to_string()),
            stmt::Return(expr) => format!("{}return{}", current_indent, match expr {
                Some(expr) => format!(" {}", expr.to_string()),
                None => String::new()
            }),
            stmt::ImportFrom(module, names, level) => {
                format!("{}from {}{} import {}",
                    current_indent,
                    match level {
                        None => String::new(),
                        Some(level) => iter::repeat(".").take(level).collect()
                    },
                    module.to_string(),
                    {
                        let mut names_ = vec!();
                        for name in names {
                            names_.push(
                                match name.asname {
                                    None => name.name,
                                    Some(asname) => format!("{} as {}", name.name, asname)
                                }
                            )
                        }
                        names_.join(", ")
                    }
                )
            },
            stmt::Import(names) => format!("{}import {}", current_indent, {
                let mut names_ = vec!();
                for name in names {
                    names_.push(
                        match name.asname {
                            None => name.name,
                            Some(asname) => format!("{} as {}", name.name, asname)
                        }
                    )
                }
                names_.join(", ")
            }),
            stmt::Expr(expr) => format!("{}{}", current_indent, expr.to_string()),
            stmt::Break => format!("{}break", current_indent),
            stmt::Delete(targets) => format!("{}del {}", current_indent, vec_to_strings_vec(targets).join(", ")),
            stmt::Pass => format!("{}pass", current_indent),
            stmt::Continue => format!("{}continue", current_indent),
            stmt::Assert(test, msg) => {
                format!("{}assert {}{}",
                    current_indent,
                    test.to_string(),
                    match msg {
                        None => String::new(),
                        Some(msg) => format!(", {}", msg.to_string())
                    }
                )
            },
            stmt::With(items, body, async) => {
                format!("{}{}with {}:\n{}",
                    current_indent,
                    if async { "async " } else { "" },
                    vec_to_strings_vec(items).join(", "),
                    statements_to_string(indent, body)
                )
            },
            stmt::Raise(exc, cause) => {
                format!("{}raise{}{}",
                    current_indent,
                    match exc {
                        None => String::new(),
                        Some(ref exc) => format!(" {}", exc.to_string())
                    },
                    match cause {
                        None => String::new(),
                        Some(ref cause) => format!(" from {}", cause.to_string())
                    }
                )
            },
            stmt::Try(body, excepthandlers, orelse, finalbody) => {
                format!("{}try:\n{}\n{}{}\n{}",
                    current_indent,
                    statements_to_string(indent, body),
                    {
                        let mut excepthandlers_ = vec!();
                        for excepthandler in excepthandlers {
                            let excepthandler = excepthandler.to_string(indent);
                            excepthandlers_.push(excepthandler);
                        }
                        excepthandlers_.join("\n")
                    },
                    if !orelse.is_empty() {
                        format!("{}else:\n{}", current_indent, statements_to_string(indent, orelse))
                    } else {
                        String::new()
                    },
                    if !finalbody.is_empty() {
                        format!("{}finally:\n{}", current_indent, statements_to_string(indent, finalbody))
                    } else {
                        String::new()
                    }
                )
            }
        }
    }
}

#[allow(dead_code)]
pub fn dump_ast(ast: &Module) -> String {
    let mut dumped_statements = vec!();
    for statement in &ast.statements {
        let dumped_statement = statement.to_string(0);
        dumped_statements.push(dumped_statement);
    }
    dumped_statements.join("\n")
}