use crate::{
    compile::{visit::*, Ctx},
    obj::prelude::*,
    syn::{ast::prelude::*, span::*},
};
use std::collections::HashMap;

pub struct NameStack {
    next_sym: usize,
    stack: Vec<HashMap<String, NameId>>,
}

impl NameStack {
    pub fn new() -> Self {
        NameStack {
            next_sym: 0,
            stack: vec![Default::default()],
        }
    }

    pub fn push(&mut self, syms: HashMap<String, NameId>) {
        self.stack.push(syms);
    }

    pub fn push_default(&mut self) {
        self.push(Default::default());
    }

    pub fn pop(&mut self) -> Option<HashMap<String, NameId>> {
        self.stack.pop()
    }

    pub fn add(&mut self, name: String) -> NameId {
        let next_sym = NameId::new(self.next_sym);
        let sym = *self.locals_mut().entry(name).or_insert(next_sym);
        if sym.id() == self.next_sym {
            self.next_sym += 1;
        }
        sym
    }

    /// Searches the top level of the symbol stack for the symbol with the supplied name.
    pub fn get_local(&self, name: &str) -> Option<NameId> {
        self.locals().get(name).copied()
    }

    /// Searches the stack from top to bottom for the symbol with the given name.
    pub fn get_scoped(&self, name: &str) -> Option<NameId> {
        self.stack
            .iter()
            .rev()
            .filter_map(|locals| locals.get(name))
            .next()
            .copied()
    }

    pub fn root(&self) -> &HashMap<String, NameId> {
        self.stack.first().expect("no root name map")
    }

    pub fn root_mut(&mut self) -> &mut HashMap<String, NameId> {
        self.stack.first_mut().expect("no root name map")
    }

    pub fn locals(&self) -> &HashMap<String, NameId> {
        self.stack.last().expect("no local name map")
    }

    pub fn locals_mut(&mut self) -> &mut HashMap<String, NameId> {
        self.stack.last_mut().expect("no local name map")
    }
}

/// Collect local names and push them to the top layer of the name stack.
pub struct CollectNames<'c, 't> {
    ctx: &'c mut Ctx,
    text: &'t str,
}

default_visitor!(Expr for CollectNames<'_, '_> where Out = ());
empty_visitor!(FunCallExpr for CollectNames<'_, '_>);
empty_visitor!(FunExpr for CollectNames<'_, '_>);

impl<'c, 't> CollectNames<'c, 't> {
    pub fn new(ctx: &'c mut Ctx, text: &'t str) -> Self {
        CollectNames { ctx, text }
    }

    pub fn collect(&mut self, stmts: &Vec<Stmt>) {
        // Collect all LHS assignments
        for stmt in stmts.iter() {
            if let Stmt::Assign(stmt) = stmt {
                self.visit(stmt);
            }
        }
    }
}

impl Visit<AssignStmt> for CollectNames<'_, '_> {
    type Out = ();

    fn visit(&mut self, stmt: &AssignStmt) -> Self::Out {
        self.visit(&stmt.lhs);
    }
}

impl Visit<BinExpr> for CollectNames<'_, '_> {
    type Out = ();

    fn visit(&mut self, expr: &BinExpr) -> Self::Out {
        self.visit(&expr.lhs);
    }
}

impl Visit<UnExpr> for CollectNames<'_, '_> {
    type Out = ();

    fn visit(&mut self, expr: &UnExpr) -> Self::Out {
        self.visit(&expr.expr)
    }
}

impl Visit<IndexExpr> for CollectNames<'_, '_> {
    type Out = ();

    fn visit(&mut self, expr: &IndexExpr) -> Self::Out {
        self.visit(&expr.expr);
    }
}

impl Visit<BaseExpr> for CollectNames<'_, '_> {
    type Out = ();

    fn visit(&mut self, expr: &BaseExpr) -> Self::Out {
        // This is a LHS standalone expr
        if let BaseExpr { kind: BaseExprKind::Ident, .. } = expr {
            let name = expr.text_at(self.text).to_string();
            self.ctx.name_stack_mut().add(name);
        }
    }
}