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

use std::iter;

trait to_string_able {
    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",
        }
    }
}

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

impl to_string_able 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 to_string_able 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 to_string_able 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!();

                // XXX: wrong order!
                for op in ops {
                    arguments.push(op.to_string().to_string())
                }
                for comparator in comparators {
                    arguments.push(comparator.to_string())
                }
                /*
                for (op, comparator) in ops.zip(comparators) {
                    let op = op.unwrap();
                    let comparator = comparator.unwrap();
                    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),
            expr::NameConstant(name) => format!("{}", name),
            expr::Attribute(value, attr, ctx) => format!("{}.{}", value.to_string(), attr),
            expr::Name(name, ctx) => format!("{}", name),
            expr::List(elts, ctx) => 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, ctx) => format!("({})", vec_to_strings_vec(elts).join(", ")),
            expr::Ellipsis => format!("..."),
        }
    }
}

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

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

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{}", current_indent, name, vec_to_strings_vec(bases).join(", "), statements_to_string(indent, body)),
            stmt::FunctionDef(name, arguments, body, decorator_list, returns) => format!("{}def {}({}):\n{}", current_indent, name, arguments.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) => format!("{}for {} in {}:\n{}", current_indent, 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 => "".to_string()
            }),
            stmt::ImportFrom(module, names, level) => {
                format!("{}from {}{} import {}",
                    current_indent,
                    {
                        match level {
                            None => String::new(),
                            Some(level) => { let dots: String = iter::repeat(".").take(level as usize).collect(); dots }
                        }
                    },
                    module.to_string(),
                    {
                        let mut names_ = vec!();
                        for name in names {
                            match name.asname {
                                None => names_.push(name.name),
                                Some(asname) => names_.push(format!("{} as {}", name.name, asname))
                            }
                        }
                        names_.join(", ")
                    }
                )
            },
            stmt::Import(names) => format!("{}import {}", current_indent, {
                let mut names_ = vec!();
                for name in names {
                    match name.asname {
                        None => names_.push(name.name),
                        Some(asname) => names_.push(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 {
                            Some(msg) => format!(", {}", msg.to_string()),
                            None => format!("")
                        }
                    }
                )
            }
        }
    }
}

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