use python_ast::{Module, Statement, Expr, BinOp, UnaryOp};

use std::iter;

impl UnaryOp {
    fn to_string(&self) -> &'static str {
        match *self {
            UnaryOp::UAdd => "+",
            UnaryOp::USub => "-",
            UnaryOp::Error => "BinOp::Error"
        }
    }
}

impl BinOp {
    fn to_string(&self) -> &'static str {
        match *self {
            BinOp::BinAdd => "+",
            BinOp::BinMult => "*",
            BinOp::BinEq => "==",
            BinOp::BinLt => "<",
            BinOp::BinGt => ">",
            BinOp::Sub => "-",
            BinOp::Div => "/",
            BinOp::Error => "BinOp::Error"
        }
    }
}

impl Expr {
    fn to_string(&self) -> String {
        match self.clone() {
            Expr::UnaryOp(op, operand) => format!("{}{}", op.to_string(), operand.to_string()),
            Expr::BinOp(a, op, b) => format!("{} {} {}", a.to_string(), op.to_string(), b.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) => format!("{}({})", func.to_string(), {
                let mut arguments = vec!();
                for arg in args {
                    arguments.push(arg.to_string());
                }
                arguments.join(", ")
            }),
            Expr::Alias(name, asname) => {
                match asname {
                    None => format!("{}", name),
                    Some(asname) => format!("{} as {}", name, asname)
                }
            }
            Expr::Attribute(lhs, rhs) => format!("{}.{}", lhs.to_string(), rhs),
            Expr::Name(name) => format!("{}", name),
            Expr::NameConstant(name) => format!("{}", name),
            Expr::Str(s) => format!("\"{}\"", s),
            Expr::Num(n) => format!("{}", n),
            Expr::List(elts) => format!("[{}]", {
                let mut elements = vec!();
                for elt in elts {
                    elements.push(elt.to_string());
                }
                elements.join(", ")
            }),
            Expr::Error => "Expr::Error".to_string()
        }
    }
}

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

impl Statement {
    fn to_string(&self, indent: usize) -> String {
        match self.clone() {
            Statement::ClassDef(name, classes, body) => format!("{}def {}({}):\n{}", make_indent(indent), name, {
                let mut bases = vec!();
                for class in classes {
                    bases.push(class.to_string());
                }
                bases.join(", ")
            }, {
                let mut statements = vec!();
                for statement in body {
                    statements.push(statement.to_string(indent + 1));
                }
                statements.join("\n")
            }),
            Statement::FunctionDef(name, arguments, body) => format!("{}def {}({}):\n{}", make_indent(indent), name, {
                let mut args = vec!();
                for arg in arguments {
                    args.push(arg.to_string());
                }
                args.join(", ")
            }, {
                let mut statements = vec!();
                for statement in body {
                    statements.push(statement.to_string(indent + 1));
                }
                statements.join("\n")
            }),
            Statement::Global(names) => format!("{}global {}", make_indent(indent), names.join(", ")),
            Statement::If(test, body, orelse) => format!("{}if {}:\n{}", make_indent(indent), test.to_string(), {
                let mut statements = vec!();
                for arg in body {
                    statements.push(arg.to_string(indent + 1));
                }

                let mut orelse_ = vec!();
                for arg in orelse {
                    orelse_.push(arg.to_string(indent + 1));
                }

                if orelse_.is_empty() {
                    statements.join("\n")
                } else {
                    format!("{}\n{}else:\n{}", statements.join("\n"), make_indent(indent), orelse_.join("\n"))
                }
            }),
            Statement::While(test, body, orelse) => format!("{}while {}:\n{}", make_indent(indent), test.to_string(), {
                let mut statements = vec!();
                for arg in body {
                    statements.push(arg.to_string(indent + 1));
                }

                let mut orelse_ = vec!();
                for arg in orelse {
                    orelse_.push(arg.to_string(indent + 1));
                }

                if orelse_.is_empty() {
                    statements.join("\n")
                } else {
                    format!("{}\n{}else:\n{}", statements.join("\n"), make_indent(indent), orelse_.join("\n"))
                }
            }),
            Statement::For(target, iter, body, orelse) => format!("{}for {} in {}:\n{}", make_indent(indent), target.to_string(), iter.to_string(), {
                let mut statements = vec!();
                for arg in body {
                    statements.push(arg.to_string(indent + 1));
                }

                let mut orelse_ = vec!();
                for arg in orelse {
                    orelse_.push(arg.to_string(indent + 1));
                }

                if orelse_.is_empty() {
                    statements.join("\n")
                } else {
                    format!("{}\n{}else:\n{}", statements.join("\n"), make_indent(indent), orelse_.join("\n"))
                }
            }),
            Statement::Assign(targets, value) => format!("{}{} = {}", make_indent(indent), {
                let mut exprs = vec!();
                for target in targets {
                    exprs.push(target.to_string());
                }
                exprs.join(", ")
            }, value.to_string()),
            Statement::AugAssign(target, op, value) => format!("{}{} {}= {}", make_indent(indent), target.to_string(), op.to_string(), value.to_string()),
            Statement::Return(expr) => format!("{}return {}", make_indent(indent), expr.to_string()),
            Statement::ImportFrom(module, names) => format!("{}from {} import {}", make_indent(indent), module.to_string(), {
                let mut exprs = vec!();
                for alias in names.iter() {
                    let alias = alias.to_string();
                    exprs.push(alias);
                }
                exprs.join(", ")
            }),
            Statement::Expr(expr) => format!("{}{}", make_indent(indent), expr.to_string()),
            Statement::Break => format!("{}break", make_indent(indent)),
            Statement::Error => format!("{}Statement::Error", make_indent(indent))
        }
    }
}

#[allow(dead_code)]
pub fn dump_ast(ast: &Module) {
    for statement in &ast.statements {
        println!("{}", statement.to_string(0));
    }
}