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

////////////////////////////////////////////////////////////////////////////////
// NameStack
////////////////////////////////////////////////////////////////////////////////

#[derive(Debug, Clone)]
pub struct NameStack {
    next_sym: usize,
    stack: Vec<HashMap<String, NameId>>,
}

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

impl NameStack {
    pub fn new() -> Self {
        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")
    }
}

////////////////////////////////////////////////////////////////////////////////
// CollectNames
////////////////////////////////////////////////////////////////////////////////

/// Collect local stack 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);
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// CollectSyms
////////////////////////////////////////////////////////////////////////////////

/// Collect symbols (identifiers and symbols included).
///
/// These are used for indexing values in attribute maps.
pub struct CollectSyms<'c, 't> {
    ctx: &'c mut Ctx,
    text: &'t str,
}

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

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

default_visitor!(Stmt for CollectSyms<'_, '_> where Out = ());
default_visitor!(Expr for CollectSyms<'_, '_> where Out = ());

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

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

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

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

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

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

impl Visit<FunCallExpr> for CollectSyms<'_, '_> {
    type Out = ();

    fn visit(&mut self, expr: &FunCallExpr) -> Self::Out {
        self.visit(&expr.expr);
        expr.args.iter().for_each(|expr| self.visit(expr));
    }
}

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

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

impl Visit<FunExpr> for CollectSyms<'_, '_> {
    type Out = ();

    fn visit(&mut self, expr: &FunExpr) -> Self::Out {
        expr.params.iter().for_each(|name| {
            self.ctx.add_sym(name.clone());
        });
        self.visit(&expr.expr);
    }
}

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

    fn visit(&mut self, expr: &BaseExpr) -> Self::Out {
        match &expr.kind {
            BaseExprKind::Ident => {
                let name = expr.text_at(self.text).to_string();
                self.ctx.add_sym(name);
            }
            BaseExprKind::Sym => {
                let name = expr.text_at(self.text).chars().skip(1).collect::<String>();
                self.ctx.add_sym(name);
            }
            BaseExprKind::List(l) | BaseExprKind::Tuple(l) => {
                l.iter().for_each(|expr| self.visit(expr))
            }
            BaseExprKind::Object(o) => o.iter().for_each(|(k, v)| {
                self.visit(k);
                self.visit(v);
            }),
            BaseExprKind::Block(stmts) => {
                stmts.iter().for_each(|stmt| self.visit(stmt));
            }
            _ => { /* no-op */ }
        }
    }
}