use rlox2_frontend::parser::{Expr, Stmt};
use rustc_hash::FxHashMap;

use crate::Runtime;
use crate::{LoxError, ResolverError};

/*====================================================================================================================*/

type ResolverScope = FxHashMap<Box<str>, ResolveStatus>;
type ResolverFrame = Vec<ResolverScope>;

type ResolverResult<T> = Result<T, ResolverError>;

/*====================================================================================================================*/

#[derive(Debug, Clone, Copy)]
enum ResolveStatus {
    Declared,
    Defined,
}

/*====================================================================================================================*/

pub fn resolve(statement: &mut Stmt, runtime: &mut Runtime) -> Result<(), LoxError> {
    let global_names = runtime.globals().keys().cloned();

    let mut resolver = Resolver::new(global_names);

    resolver.resolve_stmt(statement)?;

    Ok(())
}

#[derive(Debug)]
struct Resolver {
    global_scope: ResolverScope,
    // local_scopes: Vec<ResolverScope>,
    frames: Vec<ResolverFrame>,

    closure_vars: FxHashMap<Box<str>, usize>,
}

impl Resolver {
    fn new(global_names: impl Iterator<Item = Box<str>>) -> Self {
        let mut global_scope = ResolverScope::default();

        for name in global_names {
            global_scope.insert(name, ResolveStatus::Defined);
        }

        Resolver {
            global_scope,
            frames: vec![ResolverFrame::new()],
            closure_vars: FxHashMap::default(),
        }
    }

    fn local_scopes(&self) -> &ResolverFrame {
        self.frames.last().unwrap()
    }

    fn local_scopes_mut(&mut self) -> &mut ResolverFrame {
        self.frames.last_mut().unwrap()
    }

    fn enter_scope(&mut self) {
        self.local_scopes_mut().push(ResolverScope::default());
    }

    fn exit_scope(&mut self) {
        self.local_scopes_mut()
            .pop()
            .expect("Tried to pop from empty ResolverFrame");
    }

    fn push_frame(&mut self) {
        self.frames.push(vec![ResolverScope::default()]);
    }

    fn pop_frame(&mut self) {
        self.frames
            .pop()
            .expect("Tried to pop non-existant ResolverFrame");
    }

    fn declare(&mut self, name: impl Into<Box<str>>) -> ResolverResult<()> {
        let name = name.into();

        if let Some(local_scope) = self.local_scopes_mut().last_mut() {
            if local_scope.get(&name).is_some() {
                return Err(ResolverError::LocalMulipleDeclarations { name });
            }

            local_scope.insert(name, ResolveStatus::Declared);
        } else {
            // global variables can be declared multiple times
            self.global_scope.insert(name, ResolveStatus::Declared);
        }

        Ok(())
    }

    fn define(&mut self, name: impl Into<Box<str>>) {
        let name = name.into();

        if let Some(local_scope) = self.local_scopes_mut().last_mut() {
            local_scope.insert(name, ResolveStatus::Defined);
        } else {
            self.global_scope.insert(name, ResolveStatus::Declared);
        }
    }

    fn check(&self, name: &str) -> Option<ResolveStatus> {
        if let Some(local_scope) = self.local_scopes().last() {
            local_scope.get(name).cloned()
        } else {
            self.global_scope.get(name).cloned()
        }
    }

    fn resolve_var(&mut self, name: &str) -> ResolverResult<Expr> {
        let mut level = 0;

        for scope in self.local_scopes().iter().rev() {
            if scope.contains_key(name) {
                return Ok(Expr::local_variable(name, level));
            }
            level += 1;
        }

        for frame in self.frames.iter().rev().skip(1) {
            for scope in frame.iter().rev() {
                if scope.contains_key(name) {
                    if !self.closure_vars.contains_key(name) {
                        // the level at which the closed-over variable will be collected from
                        // from the perspective of the parameter/closure scope i.e. the outmost of the function
                        self.closure_vars
                            .insert(name.into(), level - self.local_scopes().len());
                    }

                    // distance from actual variable refernce to parameter/closure scope
                    let level = self.local_scopes().len() - 1;

                    return Ok(Expr::local_variable(name, level));
                }
                level += 1;
            }
        }

        if self.global_scope.contains_key(name) {
            return Ok(Expr::global_variable(name));
        }

        if name == "this" {
            Err(ResolverError::ThisOutsideMethod)
        } else if name == "super" {
            Err(ResolverError::SuperOutsideMethod)
        } else {
            let name = name.into();
            Err(ResolverError::UnresolvableVariable { name })
        }
    }

    fn resolve_stmt(&mut self, stmt: &mut Stmt) -> ResolverResult<()> {
        match stmt {
            Stmt::Print { expr } => self.resolve_expr(expr),
            Stmt::IfStmt {
                condition,
                then_branch,
                else_branch,
            } => {
                self.resolve_expr(condition)?;

                self.resolve_stmt(then_branch)?;

                if let Some(else_branch) = else_branch {
                    self.resolve_stmt(else_branch)?;
                }

                Ok(())
            }
            Stmt::While { condition, body } => {
                self.resolve_expr(condition)?;

                self.resolve_stmt(body)?;

                Ok(())
            }
            Stmt::VarDecl { name, initializer } => {
                self.declare(name.clone())?;

                self.resolve_expr(initializer)?;

                self.define(name.clone());

                Ok(())
            }
            Stmt::Block { statements } => {
                self.enter_scope();

                for statement in statements.iter_mut() {
                    self.resolve_stmt(statement)?;
                }

                self.exit_scope();

                Ok(())
            }
            Stmt::ExprStmt { expr } => self.resolve_expr(expr),
            Stmt::Break => Ok(()),
            Stmt::Return { expr } => self.resolve_expr(expr),
        }
    }

    fn resolve_expr(&mut self, expr: &mut Expr) -> ResolverResult<()> {
        match expr {
            Expr::Literal { literal: _ } => Ok(()),
            Expr::Unary { op: _, expr } => self.resolve_expr(expr),
            Expr::Binary { left, op: _, right } => {
                self.resolve_expr(left)?;

                self.resolve_expr(right)?;

                Ok(())
            }
            Expr::Logical { left, op: _, right } => {
                self.resolve_expr(left)?;

                self.resolve_expr(right)?;

                Ok(())
            }
            Expr::Grouping { expr } => self.resolve_expr(expr),
            Expr::Variable { name } => {
                if let Some(ResolveStatus::Declared) = self.check(name) {
                    let name = name.clone();
                    return Err(ResolverError::VarInOwnInitializer { name });
                }

                *expr = self.resolve_var(name)?;

                Ok(())
            }
            Expr::LocalVariable { name, level: _ } => {
                panic!("Tried to resolve variable {name} twice")
            }
            Expr::GlobalVariable { name } => {
                panic!("Tried to resolve variable {name} twice")
            }

            Expr::Assignment { target, value } => {
                self.resolve_expr(value)?;

                let target = target.as_mut();

                if let Expr::Variable { name } = target {
                    *target = self.resolve_var(name)?;
                } else {
                    panic!("Invalid assignment target {target}");
                }

                Ok(())
            }
            Expr::Call { callee, args } => {
                self.resolve_expr(callee)?;

                for arg in args.iter_mut() {
                    self.resolve_expr(arg)?;
                }

                Ok(())
            }
            Expr::Get { target, name: _ } => {
                self.resolve_expr(target)?;

                Ok(())
            }
            Expr::Set {
                target,
                name: _,
                value,
            } => {
                self.resolve_expr(target)?;

                self.resolve_expr(value)?;

                Ok(())
            }
            Expr::Function {
                name,
                param_names,
                closure_vars,
                body,
            } => {
                let old_closure_names = std::mem::take(&mut self.closure_vars);

                self.push_frame();

                self.declare(name.clone())?;

                self.define(name.clone());

                for param_name in param_names.iter() {
                    self.declare(param_name.clone())?;
                    self.define(param_name.clone());
                }

                self.resolve_stmt(body)?;

                self.pop_frame();

                let closure_names = std::mem::replace(&mut self.closure_vars, old_closure_names);

                for closure_var in closure_names {
                    closure_vars.push(closure_var);
                }

                Ok(())
            }
            Expr::Class {
                superclass,
                name,
                methods,
            } => {
                self.declare(name.clone())?;

                if let Some(superclass) = superclass {
                    self.resolve_expr(superclass)?;
                }

                self.define(name.clone());

                // this is the scope "this" is defined in
                self.enter_scope();

                // this should never fail! we just made a new scope
                self.declare("this").unwrap();

                if superclass.is_some() {
                    // this should never fail either! same as `this`
                    self.declare("super").unwrap();
                }

                for method in methods.iter_mut() {
                    self.resolve_expr(method)?;
                }

                self.exit_scope();

                Ok(())
            }
            Expr::This => {
                *expr = self.resolve_var("this")?;
                Ok(())
            }
            Expr::Super {
                super_var,
                this_var,
                method: _,
            } => {
                self.resolve_expr(super_var)?;

                self.resolve_expr(this_var)
            }
        }
    }
}