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

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct NameId(usize);

impl NameId {
    pub fn new(id: usize) -> Self {
        NameId(id)
    }

    pub fn id(&self) -> usize {
        self.0
    }
}

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: 'c> {
    ctx: &'c mut Ctx<'t>,
}

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

    pub fn collect(&mut self, stmts: &Vec<Stmt>) {
        for stmt in stmts.iter() {
            self.visit_stmt(stmt);
        }
    }

    fn visit_base_expr(&mut self, expr: &BaseExpr) {
        match &expr.kind {
            BaseExprKind::Ident => {
                let name = expr.text_at(self.ctx.text()).to_string();
                self.ctx.name_stack_mut().add(name);
            }
            BaseExprKind::List(l) | BaseExprKind::Tuple(l) => {
                l.iter().for_each(|expr| self.visit_expr(expr));
            }
            BaseExprKind::Object(o) => {
                o.iter().for_each(|(key, value)| {
                    self.visit_expr(key);
                    self.visit_expr(value);
                });
            }
            _ => { /* no-op */ }
        }
    }
}

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

    fn visit_expr(&mut self, expr: &Expr) -> Self::Out {
        match expr {
            Expr::Bin(b) => {
                self.visit_expr(&b.lhs);
                // don't descend into dot-expressions - these are not names, but symbols
                if b.op != BinOp::Dot {
                    self.visit_expr(&b.rhs);
                }
            }
            Expr::Un(u) => {
                self.visit_expr(&u.expr);
            }
            Expr::FunCall(f) => {
                self.visit_expr(&f.expr);
                f.args.iter().for_each(|expr| self.visit_expr(expr));
            }
            Expr::Index(i) => {
                self.visit_expr(&i.expr);
                self.visit_expr(&i.index);
            }
            Expr::Fun(_) => { /* no-op */ }
            Expr::Base(b) => self.visit_base_expr(b),
        }
    }

    fn visit_stmt(&mut self, stmt: &Stmt) -> Self::Out {
        match stmt {
            Stmt::Assign(AssignStmt { lhs, rhs, .. }) => {
                self.visit_expr(lhs);
                self.visit_expr(rhs);
            }
            Stmt::Expr(expr) => {
                self.visit_expr(expr);
            }
        }
    }
}