use std::collections::HashMap;

use crate::{collect_lines, error::{ErrorKind, KabelError}, out_of_scope, out_of_scope_var, parser::{ASTType, Lit, Name, AST}};

pub struct Analyzer {
    text: Vec<String>,
    symbol_table: Vec<HashMap<String, Symbol>>,
    pub errors: Vec<KabelError>,
}

impl Analyzer {
    pub fn new(text: String) -> Self {
        Self {
            text: text.lines().collect::<Vec<&str>>().iter().map(|s| s.to_string()).collect(),
            symbol_table: vec![HashMap::new()],
            errors: Vec::new(),
        }
    }
    pub fn visit(&mut self, ast: AST) {
        use ASTType::*;
        match ast.ast_type {
            Program(asts) => {
                self.visit_program(asts);
            }
            Function(name, args, block) => {
                self.visit_function(name, args, *block);
            }
            Return(expr) => {
                self.visit_return(*expr);
            }
            Loop(block) => {
                self.visit_loop(*block);
            }
            While(condition, block) => {
                self.visit_while(*condition, *block);
            }
            For(expr1, expr2, expr3, block) => {
                self.visit_for(*expr1, *expr2, *expr3, *block);
            }
            If(condition, block, else_expr) => {
                self.visit_if(*condition, *block, *else_expr);
            }
            Block(stmts) => {
                self.visit_block(stmts);
            }
            // REMOVE LATER
            Print(expr) => {
                self.visit_print(*expr);
            }
            Decl(name, expr) => {
                self.visit_decl(name, *expr);
            }
            Assign(ref name, ref expr) => {
                self.visit_assign(ast.clone(), name.clone(), *expr.clone());
            }
            Ternary(condition, true_expr, false_expr) => {
                self.visit_ternary(*condition, *true_expr, *false_expr);
            }
            Subscript(array, index) => {
                self.visit_subscript(*array, *index);
            }
            Binary(left, _oper, right) => {
                self.visit_binary(*left, *right);
            }
            Unary(_oper, right) => {
                self.visit_unary(*right);
            }
            Lit(ref lit) => {
                self.visit_lit(ast.clone(), lit.clone());
            }
            Call(ref ident, ref args) => {
                self.visit_call(ast.clone(), ident.clone(), args.clone());
            }
            /*Member(left, right) => {
                self.visit_member(*left, *right);
            }*/
            _ => {} // not implemented
        }
    }
    pub fn visit_program(&mut self, asts: Vec<AST>) {
        for ast in asts {
            self.visit(ast);
        }
    }
    pub fn visit_function(&mut self, name: Name, args: Vec<Name>, block: AST) {
        self.symbol_table.last_mut().unwrap().insert(name.name.clone(), Symbol::Function(args.len()));
        self.symbol_table.push(HashMap::new());
        for arg in args {
            self.symbol_table.last_mut().unwrap().insert(arg.name, Symbol::Var);
        }
        self.visit(block);
        self.symbol_table.pop();
    }
    pub fn visit_return(&mut self, expr: Option<AST>) {
        if let Some(expr) = expr {
            self.visit(expr);
        }
    }
    pub fn visit_loop(&mut self, block: AST) {
        self.visit(block);
    }
    pub fn visit_while(&mut self, condition: AST, block: AST) {
        self.visit(condition);
        self.visit(block);
    }
    pub fn visit_for(&mut self, expr1: Option<AST>, expr2: Option<AST>, expr3: Option<AST>, block: AST) {
        if let Some(expr) = expr1 {
            self.visit(expr);
        }
        if let Some(expr) = expr2 {
            self.visit(expr);
        }
        if let Some(expr) = expr3 {
            self.visit(expr);
        }
        self.visit(block);
    }
    pub fn visit_if(&mut self, condition: AST, block: AST, else_expr: Option<AST>) {
        self.visit(condition);
        self.visit(block);
        if let Some(else_expr) = else_expr {
            self.visit(else_expr);
        }
    }
    pub fn visit_block(&mut self, stmts: Vec<AST>) {
        self.symbol_table.push(HashMap::new());
        for stmt in stmts {
            self.visit(stmt);
        }
        self.symbol_table.pop();
    }
    // REMOVE LATER
    pub fn visit_print(&mut self, expr: AST) {
        self.visit(expr);
    }
    pub fn visit_decl(&mut self, name: Name, expr: AST) {
        self.visit(expr);
        self.symbol_table.last_mut().unwrap().insert(name.name, Symbol::Var);
    }
    pub fn visit_assign(&mut self, ast: AST, name: Name, expr: AST) {
        self.visit(expr.clone());
        if !self.symbol_table.last().unwrap().contains_key(&name.name) {
            self.errors.push(out_of_scope_var!(self, "Variable \"{}\" not in scope", name, ast))
        }
    }
    pub fn visit_ternary(&mut self, condition: AST, true_expr: AST, false_expr: AST) {
        self.visit(condition);
        self.visit(true_expr);
        self.visit(false_expr);
    }
    pub fn visit_subscript(&mut self, array: AST, index: AST) {
        self.visit(array);
        self.visit(index);
    }
    pub fn visit_binary(&mut self, left: AST, right: AST) {
        self.visit(left);
        self.visit(right);
    }
    pub fn visit_unary(&mut self, right: AST) {
        self.visit(right);
    }
    pub fn visit_lit(&mut self, ast: AST, lit: Lit) {
        match lit {
            Lit::Ident(name) => {
                if !self.resolve_var(&name) {
                    self.errors.push(out_of_scope!(self, "Variable \"{}\" not in scope", name, ast))
                }
            }
            _ => {}
        }
    }
    pub fn visit_call(&mut self, ast: AST, ident: Name, args: Vec<AST>) {
        if !self.resolve_function(&ast, &ident.name, args.len()) {
            self.errors.push(out_of_scope!(self, "Function \"{}\" not in scope", ident.name, ast))
        }
        for arg in args {
            self.visit(arg);
        }
    }
    // TODO: make visit_member not throw out of scope errors
    /*pub fn visit_member(&mut self, left: AST, right: AST) {
        self.visit(left);
        self.visit(right);
    }*/
    fn resolve_var(&self, name: &String) -> bool {
        for scope in self.symbol_table.iter().rev() {
            if matches!(scope.get(name), Some(Symbol::Var)) {
                return true;
            }
        }
        false
    }
    fn resolve_function(&mut self, ast: &AST, name: &String, arity: usize) -> bool {
        for scope in self.symbol_table.iter().rev() {
            if let Some(Symbol::Function(f_arity)) = scope.get(name) {
                if *f_arity == arity {
                    return true;
                } else {
                    self.errors.push(
                        KabelError::new(
                            ErrorKind::OutOfScope,
                            format!("Function {} has {} argument, provided {}", name, *f_arity, arity),
                            ast.start_line,
                            ast.start_column,
                            collect_lines!(self.text[ast.start_line-1..ast.end_line-1]),
                        )
                    );
                    return true;
                }
            }
        }
        false
    }
}

pub enum Symbol {
    Var,
    Function(usize),
}