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

use std::collections::HashMap;

#[derive(Clone, Debug, PartialEq, Eq)]
enum Type {
    Top(usize),
    Unit,
    Bool,
    Int,
    Str,
    List(Box<Type>),
    Function(Vec<Type>, Vec<Type>),
    Bottom
}

struct TypeVariable {
    id: usize,
    instance: Box<TypeVariable>,
}

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 Typing {
    environment: Vec<HashMap<String, Type>>,
    next_id: usize,
}

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

    fn visit_statement(&mut self, statement: Statement) -> Type {
        match statement {
            Statement::ClassDef(name, classes, body) => {
                Type::Bottom
            },
            Statement::FunctionDef(name, arguments, body) => {
                let mut env = self.environment.pop().unwrap();
                self.environment.push(env.clone());

                let nb_args = arguments.len();
                for expr in arguments {
                    let type_ = self.visit_expr(expr.clone());
                    let name = match expr {
                        Expr::Name(name) => name,
                        _ => panic!()
                    };
                    env.insert(name.clone(), type_.clone());
                }
                self.environment.push(env);

                for statement in body {
                    self.visit_statement(statement);
                }
                self.environment.pop();

                let mut types = Vec::with_capacity(nb_args);
                for _ in 0..nb_args {
                    self.next_id += 1;
                    types.push(Type::Top(self.next_id));
                }
                Type::Bottom
            },
            Statement::Global(_) => {
                Type::Bottom
            },
            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);
                }
                Type::Bottom
            },
            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);
                }
                Type::Bottom
            },
            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);
                }
                Type::Bottom
            },
            Statement::Assign(targets, value) => {
                let type_ = self.visit_expr(value);
                if targets.len() != 1 {
                    panic!();
                }
                println!("{:?}", self.environment.clone());
                let mut env = self.environment.pop().unwrap();
                for target in targets {
                    let name = match target {
                        Expr::Name(name) => name,
                        _ => panic!()
                    };
                    println!("{} => {:?}", name, type_);
                    env.insert(name, type_.clone());
                }
                self.environment.push(env);
                println!("{:?}", self.environment.clone());
                Type::Bottom
            },
            Statement::AugAssign(target, op, value) => {
                let value = self.visit_expr(value);
                let target = self.visit_expr(target);
                Type::Bottom
            },
            Statement::Return(expr) => {
                self.visit_expr(expr)
            },
            Statement::ImportFrom(module, names) => {
                for alias in names {
                    self.visit_expr(alias);
                }
                Type::Bottom
            },
            Statement::Expr(expr) => {
                self.visit_expr(expr)
            },
            Statement::Break => {
                Type::Bottom
            },
            Statement::Error => {
                println!("Statement::Error");
                panic!()
            },
        }
    }

    fn visit_expr(&mut self, expr: Expr) -> Type {
        let expr_str = format!("{:?}", expr);
        let type_ = match expr {
            Expr::UnaryOp(op, operand) => {
                let type_operand = self.visit_expr(*operand);
                type_operand
            },
            Expr::BinOp(left, op, right) => {
                let type_left = self.visit_expr(*left);
                let type_right = self.visit_expr(*right);

                if op == BinOp::BinMult && type_left == Type::Str && type_right == Type::Int {
                    return Type::Str;
                }

                if type_left != type_right {
                    return Type::Bottom;
                }
                if op == BinOp::BinAdd || op == BinOp::BinMult {
                    if type_left == Type::Int {
                        Type::Int
                    } else if type_left == Type::Str {
                        Type::Str
                    } else {
                        Type::Bottom
                    }
                } else if op == BinOp::BinEq || op == BinOp::BinLt {
                    Type::Bool
                } else {
                    Type::Bottom
                }
            },
            Expr::Compare(_, _, _) => Type::Bool,
            Expr::Call(func, args) => {
                let func = *func;
                let func = match func {
                    Expr::Name(arg) => arg,
                    _ => panic!()
                };
                let mut types = vec!();
                for expr in args {
                    let type_ = self.visit_expr(expr);
                    types.push(type_);
                }
                println!("types: {:?}", types);
                let env = match self.environment.pop() {
                    Some(env) => env,
                    None => return Type::Bottom
                };
                if env.contains_key(&func) {
                    let value = env.get(&func).unwrap().clone();
                    self.environment.push(env);
                    value
                } else {
                    self.environment.push(env);
                    Type::Bottom
                }
            },
            Expr::Alias(_, _) => Type::Bottom,
            Expr::Attribute(lhs, rhs) => {
                let type_left = self.visit_expr(*lhs);
                println!("coucou {:?}", type_left);
                Type::Bottom
            },
            Expr::Name(id) => {
                let env = match self.environment.pop() {
                    Some(env) => env,
                    None => return Type::Bottom
                };
                if env.contains_key(&id) {
                    let value = env.get(&id).unwrap().clone();
                    self.environment.push(env);
                    value
                } else {
                    self.environment.push(env);
                    self.next_id += 1;
                    Type::Top(self.next_id)
                }
            },
            Expr::NameConstant(value) => {
                if value == "True" || value == "False" {
                    Type::Bool
                } else if value == "None" {
                    Type::Unit
                } else {
                    Type::Bottom
                }
            },
            Expr::Str(_) => Type::Str,
            Expr::Num(_) => Type::Int,
            Expr::List(_) => Type::List(Box::new(Type::Bottom)),
            Expr::Error => {
                println!("Expr::Error");
                panic!()
            }
        };
        println!("{} => {:?}", expr_str, type_);
        type_
    }
}

#[allow(dead_code)]
pub fn type_ast(ast: Module) {
    let mut environment = HashMap::new();
    environment.insert("int".to_string(), Type::Int);
    environment.insert("__name__".to_string(), Type::Str);
    environment.insert("print".to_string(), Type::Unit);
    let mut typing = Typing{environment: vec!(environment), next_id: 0};
    typing.visit_module(ast);
}