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

trait Visitor<T> {
    fn visit_module(&mut self, module: Module) -> T;
    fn visit_statement(&mut self, statement: Statement) -> T;
    fn visit_expr(&mut self, expr: Expr) -> T;
}

struct Rewrite {
}

impl Visitor<()> for Rewrite {
    fn visit_module(&mut self, module: Module) -> () {
        for statement in module.statements {
            println!("{:?}", statement);
            self.visit_statement(statement);
        }
    }

    fn visit_statement(&mut self, statement: Statement) -> () {
        match statement {
            Statement::FunctionDef(Expr::Name(name), arguments, body) => {
                for expr in arguments {
                    self.visit_expr(expr);
                }
                for statement in body {
                    self.visit_statement(statement);
                }
            },
            Statement::FunctionDef(_, _, _) => {
                println!("Statement:FunctionDef Error");
                panic!()
            },
            Statement::If(test, body, orelse) => {
                self.visit_expr(test);
                for statement in body {
                    self.visit_statement(statement);
                }
                for statement in orelse {
                    self.visit_statement(statement);
                }
            },
            Statement::While(test, body, orelse) => {
                self.visit_expr(test);
                for statement in body {
                    self.visit_statement(statement);
                }
                for statement in orelse {
                    self.visit_statement(statement);
                }
            },
            Statement::For(target, iter, body, orelse) => {
                self.visit_expr(target);
                self.visit_expr(iter);
                for statement in body {
                    self.visit_statement(statement);
                }
                for statement in orelse {
                    self.visit_statement(statement);
                }
            },
            Statement::Assign(targets, value) => {
                self.visit_expr(value);
                for target in targets {
                    self.visit_expr(target);
                }
            },
            Statement::AugAssign(target, op, value) => {
                self.visit_expr(value);
                self.visit_expr(target);
            },
            Statement::Return(expr) => {
                self.visit_expr(expr);
            },
            Statement::ImportFrom(module, names) => {
                //self.visit_expr(module);
            },
            Statement::Expr(expr) => {
                self.visit_expr(expr);
            },
            Statement::Break => {},
            Statement::Error => {
                println!("Statement::Error");
                panic!()
            },
        }
    }

    fn visit_expr(&mut self, expr: Expr) -> () {
        match expr {
            Expr::BinOp(left, op, right) => {
                self.visit_expr(*left);
                self.visit_expr(*right);
            },
            Expr::Compare(left, ops, comparators) => {
                self.visit_expr(*left);
            },
            Expr::Call(func, args) => {
                let func = *func;
                let func = match func {
                    Expr::Name(arg) => arg,
                    _ => panic!()
                };
                for expr in args {
                    self.visit_expr(expr);
                }
            },
            Expr::Alias(_, _) => (),
            Expr::Name(_) => (),
            Expr::NameConstant(value) => (),
            Expr::Str(_) => (),
            Expr::Num(_) => (),
            Expr::Error => {
                println!("Expr::Error");
                panic!()
            }
        }
    }
}

#[allow(dead_code)]
pub fn rewrite_ast(ast: Module) {
    let mut typing = Rewrite{};
    typing.visit_module(ast);
}