use std::collections::HashMap;

use crate::{
    ast::{ASTType, LhsAssign, LhsAssignType, AST},
    ast_error, ast_from_ast,
    error::{ErrorKind, KabelError},
    extension::Extension,
    out_of_scope_var,
};

pub struct Resolver {
    text: Vec<String>,
    symbol_table: Vec<HashMap<String, (Symbol, usize)>>, // (Symbol, reference to locals)
    pub locals: Vec<Vec<usize>>,                         // scope
    pub scope: usize,
    pub errors: Vec<KabelError>,
}

impl Resolver {
    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()],
            locals: vec![Vec::new()],
            scope: 0,
            errors: Vec::new(),
        }
    }
    pub fn visit(&mut self, ast: AST) -> AST {
        use ASTType::*;
        match ast.kind {
            Program(asts) => {
                let mut program = Vec::new();
                for ast in asts {
                    let ast = self.visit(ast.clone());
                    program.push(ast)
                }
                AST {
                    kind: ASTType::Program(program),
                    extensions: Vec::new(),
                    start_line: 0,
                    end_line: 0,
                    start_column: 0,
                    end_column: 0,
                }
            }
            Function(name, args, block) => {
                if self.resolve_var(&name.name).0 {
                    self.errors.push(out_of_scope_var!(self,
                            ErrorKind::FunctionAlreadyDeclaredVariable, ast,
                            "Function \"{}\" already declared", name ;
                            "hint: has variable \"{}\" already been declared?", name.name));
                } else {
                    let resolution = self.resolve_function(&name.name, 0);
                    if resolution.is_ok() {
                        self.errors.push(out_of_scope_var!(self,
                                ErrorKind::FunctionAlreadyDeclaredFunction, ast,
                                "Function \"{}\" already declared", name ;
                                "hint: has function \"{}\" already been declared?", name.name));
                    }
                    if let Err((kind, _, _)) = resolution {
                        if kind == ErrorKind::IncorrectArity {
                            self.errors.push(out_of_scope_var!(self,
                                    ErrorKind::FunctionAlreadyDeclaredFunction, ast,
                                    "Function \"{}\" already declared", name ;
                                    "hint: has function \"{}\" already been declared?", name.name));
                        }
                    }
                }
                self.locals
                    .last_mut()
                    .expect("locals last in function push")
                    .push(self.scope);
                self.symbol_table.last_mut().unwrap().insert(
                    name.name.clone(),
                    (
                        Symbol::Function(args.len()),
                        self.locals
                            .last()
                            .expect("locals last in function symbol len")
                            .len()
                            - 1,
                    ),
                );

                self.locals.push(Vec::new());
                self.symbol_table.push(HashMap::new());
                self.locals
                    .last_mut()
                    .expect("locals last in function self-reference push")
                    .push(self.scope + 1);
                self.symbol_table.last_mut().unwrap().insert(
                    name.name.clone(),
                    (
                        Symbol::Var,
                        self.locals
                            .last()
                            .expect("locals last in function self-reference len")
                            .len()
                            - 1,
                    ),
                );
                for arg in args.clone() {
                    self.locals
                        .last_mut()
                        .expect("locals last in function arg push")
                        .push(self.scope + 1);
                    self.symbol_table.last_mut().unwrap().insert(
                        arg.name,
                        (
                            Symbol::Var,
                            self.locals
                                .last()
                                .expect("locals last in function arg len")
                                .len()
                                - 1,
                        ),
                    );
                }
                let block = self.visit(*block);
                self.symbol_table.pop();
                self.locals.pop();
                AST {
                    kind: Function(name, args, Box::new(block)),
                    extensions: vec![Extension::Resolution(
                        self.scope,
                        self.locals
                            .last()
                            .expect("locals last in function ast len")
                            .len()
                            - 1,
                    )],
                    start_line: ast.start_line,
                    end_line: ast.end_line,
                    start_column: ast.start_column,
                    end_column: ast.end_column,
                }
            }
            Return(expr) => {
                if let Some(expr) = *expr {
                    let expr = self.visit(expr);
                    return ast_from_ast!(AST, Return(Box::new(Some(expr))), ast, ast);
                }
                ast_from_ast!(AST, Return(Box::new(None)), ast, ast)
            }
            Loop(block) => {
                let block = self.visit(*block);
                ast_from_ast!(AST, Loop(Box::new(block)), ast, ast)
            }
            While(condition, block) => {
                let condition = self.visit(*condition);
                let block = self.visit(*block);
                ast_from_ast!(AST, While(Box::new(condition), Box::new(block)), ast, ast)
            }
            Break => {
                ast_from_ast!(AST, Break, ast, ast)
            }
            Continue => {
                ast_from_ast!(AST, Continue, ast, ast)
            }
            For(expr1, expr2, expr3, block) => {
                self.symbol_table.push(HashMap::new());
                self.scope += 1;
                let mut n_expr1 = None;
                let mut n_expr2 = None;
                let mut n_expr3 = None;
                if let Some(expr) = *expr1 {
                    n_expr1 = Some(self.visit(expr));
                }
                if let Some(expr) = *expr2 {
                    n_expr2 = Some(self.visit(expr));
                }
                if let Some(expr) = *expr3 {
                    n_expr3 = Some(self.visit(expr));
                }
                let block = self.visit(*block);
                while let Some(scope) = self.locals.last().expect("locals failed in For").last() {
                    if self.scope == *scope {
                        self.locals
                            .last_mut()
                            .expect("locals failed in For pop")
                            .pop();
                    } else {
                        break;
                    }
                }
                self.scope -= 1;
                self.symbol_table.pop();
                ast_from_ast!(
                    AST,
                    For(
                        Box::new(n_expr1),
                        Box::new(n_expr2),
                        Box::new(n_expr3),
                        Box::new(block)
                    ),
                    ast,
                    ast
                )
            }
            If(condition, block, else_expr) => {
                let condition = self.visit(*condition);
                let block = self.visit(*block);
                let mut n_else_expr = None;
                if let Some(else_expr) = *else_expr {
                    n_else_expr = Some(self.visit(else_expr));
                }
                ast_from_ast!(
                    AST,
                    If(Box::new(condition), Box::new(block), Box::new(n_else_expr)),
                    ast,
                    ast
                )
            }
            Block(stmts) => {
                self.symbol_table.push(HashMap::new());
                self.scope += 1;
                let mut n_stmts = Vec::new();
                for stmt in stmts {
                    n_stmts.push(self.visit(stmt));
                }
                /*for (index, scope) in self.locals.clone().iter().enumerate() {
                    if self.scope == *scope {
                        self.locals.remove(index);
                    }
                }*/
                while let Some(scope) = self.locals.last().expect("locals last in block").last() {
                    if self.scope == *scope {
                        self.locals
                            .last_mut()
                            .expect("locals last in block pop")
                            .pop();
                    } else {
                        break;
                    }
                }
                self.scope -= 1;
                self.symbol_table.pop();
                ast_from_ast!(AST, Block(n_stmts), ast, ast)
            }
            Decl(name, expr) => {
                let expr = self.visit(*expr);
                self.locals
                    .last_mut()
                    .expect("locals last in decl push")
                    .push(self.scope);
                let resolution = self.resolve_var(&name.name);
                if resolution.0 && resolution.2 == self.scope {
                    self.errors.push(out_of_scope_var!(self,
                            ErrorKind::VariableAlreadyDeclaredVariable, ast,
                            "Variable \"{}\" already declared", name ;
                            "hint: has variable \"{}\" already been declared?", name.name));
                } else {
                    let resolution = self.resolve_function(&name.name, 0);
                    if resolution.is_ok() {
                        self.errors.push(out_of_scope_var!(self,
                                ErrorKind::VariableAlreadyDeclaredFunction, ast,
                                "Variable \"{}\" already declared", name ;
                                "hint: has function \"{}\" already been declared?", name.name));
                    }
                    if let Err((kind, _, _)) = resolution {
                        if kind == ErrorKind::IncorrectArity {
                            self.errors.push(out_of_scope_var!(self,
                                    ErrorKind::VariableAlreadyDeclaredFunction, ast,
                                    "Variable \"{}\" already declared", name ;
                                    "hint: has function \"{}\" already been declared?", name.name));
                        }
                    }
                }
                self.symbol_table.last_mut().unwrap().insert(
                    name.name.clone(),
                    (
                        Symbol::Var,
                        self.locals
                            .last()
                            .expect("locals last in decl symbol")
                            .len()
                            - 1,
                    ),
                );
                AST {
                    kind: Decl(name, Box::new(expr)),
                    extensions: vec![Extension::Resolution(
                        self.scope,
                        self.locals.last().expect("locals last in decl ast").len() - 1,
                    )],
                    start_line: ast.start_line,
                    end_line: ast.end_line,
                    start_column: ast.start_column,
                    end_column: ast.end_column,
                }
            }
            Expr(expr) => {
                let expr = self.visit(*expr);
                ast_from_ast!(AST, Expr(Box::new(expr)), ast, ast)
            }
            // REMOVE LATER
            Print(expr) => {
                let expr = self.visit(*expr);
                ast_from_ast!(AST, Print(Box::new(expr)), ast, ast)
            }
            Assign(lhs, expr) => {
                let expr = self.visit(*expr);
                match lhs.kind.clone() {
                    LhsAssignType::Ident(name) => {
                        let resolution = self.resolve_var(&name);
                        if !resolution.0 {
                            self.errors.push(out_of_scope_var!(
                                self,
                                ErrorKind::OutOfScope,
                                lhs,
                                "Variable \"{}\" not in scope",
                                name
                            ));
                        }
                        AST {
                            kind: Assign(lhs, Box::new(expr)),
                            extensions: vec![Extension::Resolution(self.scope, resolution.1)],
                            start_line: ast.start_line,
                            end_line: ast.end_line,
                            start_column: ast.start_column,
                            end_column: ast.end_column,
                        }
                    }
                    LhsAssignType::Subscript(name, index) => {
                        let name = self.visit(*name);
                        let index = self.visit(*index);
                        AST {
                            kind: Assign(
                                ast_from_ast!(
                                    LhsAssign,
                                    LhsAssignType::Subscript(Box::new(name), Box::new(index)),
                                    lhs,
                                    lhs
                                ),
                                Box::new(expr),
                            ),
                            extensions: Vec::new(),
                            start_line: ast.start_line,
                            end_line: ast.end_line,
                            start_column: ast.start_column,
                            end_column: ast.end_column,
                        }
                    }
                }
            }
            Anonymous(params, block) => {
                self.locals.push(Vec::new());
                self.symbol_table.push(HashMap::new());
                self.locals
                    .last_mut()
                    .expect("locals last in anonymous self-reference push")
                    .push(self.scope);
                for param in params.clone() {
                    self.locals
                        .last_mut()
                        .expect("locals last in anonymous param push")
                        .push(self.scope + 1);
                    self.symbol_table.last_mut().unwrap().insert(
                        param.name,
                        (
                            Symbol::Var,
                            self.locals
                                .last()
                                .expect("locals last in anonymous param len")
                                .len()
                                - 1,
                        ),
                    );
                }
                let block = self.visit(*block);
                self.symbol_table.pop();
                self.locals.pop();
                ast_from_ast!(AST, Anonymous(params, Box::new(block)), ast, ast)
            }
            Ternary(condition, true_expr, false_expr) => {
                let condition = self.visit(*condition);
                let true_expr = self.visit(*true_expr);
                let false_expr = self.visit(*false_expr);
                ast_from_ast!(
                    AST,
                    Ternary(
                        Box::new(condition),
                        Box::new(true_expr),
                        Box::new(false_expr)
                    ),
                    ast,
                    ast
                )
            }
            Subscript(array, index) => {
                let array = self.visit(*array);
                let index = self.visit(*index);
                ast_from_ast!(AST, Subscript(Box::new(array), Box::new(index)), ast, ast)
            }
            Binary(left, oper, right) => {
                let left = self.visit(*left);
                let right = self.visit(*right);
                ast_from_ast!(AST, Binary(Box::new(left), oper, Box::new(right)), ast, ast)
            }
            Unary(oper, right) => {
                let right = self.visit(*right);
                ast_from_ast!(AST, Unary(oper, Box::new(right)), ast, ast)
            }
            Lit(ref lit) => {
                let lit = lit.clone();
                match lit {
                    crate::ast::Lit::Ident(ref name) => {
                        let resolution = self.resolve_var(name);
                        if !resolution.0 {
                            self.errors.push(ast_error!(
                                self,
                                ErrorKind::OutOfScope,
                                ast,
                                "Variable \"{}\" not in scope",
                                name
                            ))
                        } else {
                            return AST {
                                kind: Lit(lit),
                                extensions: vec![Extension::Resolution(self.scope, resolution.1)],
                                start_line: ast.start_line,
                                end_line: ast.end_line,
                                start_column: ast.start_column,
                                end_column: ast.end_column,
                            };
                        }
                    }
                    _ => {}
                }
                ast_from_ast!(AST, Lit(lit), ast, ast)
            }
            Call(ident, args) => {
                let ident = self.visit(*ident.clone());
                let mut n_args = Vec::new();
                for arg in args {
                    n_args.push(self.visit(arg));
                }
                return AST {
                    kind: Call(Box::new(ident), n_args),
                    extensions: Vec::new(),
                    start_line: ast.start_line,
                    end_line: ast.end_line,
                    start_column: ast.start_column,
                    end_column: ast.end_column,
                };
            }
            /*Member(left, right) => {
                self.visit_member(*left, *right);
            }*/
            _ => {
                panic!("not implemented")
            } // not implemented
        }
    }
    // 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, usize, usize) {
        for (scope_num, scope) in self.symbol_table.iter().enumerate().rev() {
            if let Some((Symbol::Var, place)) | Some((Symbol::Function(_), place)) = scope.get(name)
            {
                return (true, *place, scope_num);
            }
        }
        (false, 0, 0)
    }
    fn resolve_function(
        &mut self,
        name: &String,
        arity: usize,
    ) -> Result<usize, (ErrorKind, Option<usize>, Option<usize>)> {
        for scope in self.symbol_table.iter().rev() {
            if let Some((Symbol::Function(f_arity), place)) = scope.get(name) {
                if *f_arity == arity {
                    return Ok(*place);
                } else {
                    return Err((ErrorKind::IncorrectArity, Some(*f_arity), Some(arity)));
                }
            }
        }
        Err((ErrorKind::OutOfScope, None, None))
    }
}

#[derive(Debug, Clone, Copy)]
pub enum Symbol {
    Var,
    Function(usize),
    Anonymous(usize),
}