use crate::compile::{scope::GlobalScope, thunk::Thunk};
use crate::obj::prelude::*;
use crate::obj::ObjPtr;
use crate::syn::ast::*;
use crate::vm::inst::Inst;

pub struct Compiler {
    scope: GlobalScope,
}

impl Compiler {
    pub fn new() -> Self {
        Self {
            scope: Default::default(),
        }
    }

    /// Gets the global scope.
    pub fn scope(&self) -> &GlobalScope {
        &self.scope
    }

    /// Gets the global scope, mutably.
    pub fn scope_mut(&mut self) -> &mut GlobalScope {
        &mut self.scope
    }

    /// Compile the given list of AST statements, returning the body.
    pub fn compile_stmts(&mut self, stmts: &Vec<SpStmt>) -> Vec<Inst> {
        self.gather_names(stmts);
        let mut thunks = Vec::new();
        for stmt in stmts {
            thunks.push(self.emit_stmt(stmt));
        }
        Thunk::List(thunks).flatten()
    }

    fn emit_stmt(&mut self, stmt: &SpStmt) -> Thunk {
        match stmt.inner() {
            Stmt::Assign(lhs_expr, expr) => match lhs_expr.inner() {
                AssignLhs::Name(n) => {
                    let name = self.scope().lookup_scoped(n).unwrap();
                    let mut thunk = self.emit_expr(expr);
                    thunk.push(Inst::Store(name));
                    thunk
                }
                AssignLhs::Complex(_, _) => {
                    todo!("Complex LHS foo.bar assign")
                }
            },
            Stmt::Expr(expr) => {
                let mut thunk = self.emit_expr(expr);
                thunk.push(Inst::Pop);
                thunk
            }
            Stmt::If { .. } => todo!(),
            Stmt::Def { .. } => todo!(),
            Stmt::Import { .. } => todo!(),
        }
    }

    fn emit_expr(&mut self, expr: &SpExpr) -> Thunk {
        match expr.inner() {
            Expr::Call(expr, args) => {
                let mut thunk = Vec::with_capacity(args.len() + 1);
                thunk.push(self.emit_expr(expr));
                thunk.extend(args.iter().map(|arg| self.emit_expr(arg)));
                let mut thunk = Thunk::List(thunk);
                thunk.push(Inst::Call(args.len()));
                thunk
            }
            Expr::Get(expr, name) => {
                let mut thunk = self.emit_expr(expr);
                thunk.push(Inst::GetAttr(name.clone()));
                thunk
            }
            Expr::Atom(atom) => self.emit_atom(atom),
        }
    }

    fn emit_atom(&mut self, atom: &SpAtom) -> Thunk {
        match atom.inner() {
            Atom::Int(i) => Thunk::Block(vec![Inst::Push(ObjPtr::new(Int::new(*i)))]),
            Atom::Float(_) => todo!(), // Thunk::Block(vec![Inst::Push(Gc::new(Float::new(*f)))]),
            Atom::Str(s) => Thunk::Block(vec![Inst::Push(ObjPtr::new(Str::new(s.clone())))]),
            Atom::Sym(_) => todo!(),
            Atom::Name(n) => {
                // Look up the symbol first. It may already exist.
                // If it doesn't exist, create it. It may be created
                // dynamically.
                let name = self
                    .scope()
                    .lookup_scoped(n)
                    .unwrap_or_else(|| self.scope_mut().insert_local(n));
                Thunk::Block(vec![Inst::Load(name)])
            }
        }
    }

    /// Gather all names from the list of statements and add them to the scope.
    fn gather_names(&mut self, stmts: &Vec<SpStmt>) {
        for stmt in stmts {
            match stmt.inner() {
                Stmt::Assign(lhs, _) => match lhs.inner() {
                    AssignLhs::Name(name) => {
                        self.scope_mut().insert_local(name);
                    }
                    AssignLhs::Complex(expr, _) => self.gather_expr_names(expr.inner()),
                },
                Stmt::Def { name, .. } => {
                    self.scope_mut().insert_local(name);
                }
                Stmt::Expr(expr) => {
                    // Get the names of the variables used in this expr
                    self.gather_expr_names(expr.inner());
                }
                Stmt::If { .. } => {
                    todo!()
                }
                Stmt::Import { .. } => {
                    todo!()
                }
            }
        }
    }

    /// Gather the names used inside of an expression, recursively.
    fn gather_expr_names(&mut self, expr: &Expr) {
        match expr {
            Expr::Call(fun, args) => {
                self.gather_expr_names(fun.inner());
                for arg in args {
                    self.gather_expr_names(arg.inner());
                }
            }
            Expr::Get(head, _) => {
                self.gather_expr_names(head.inner());
            }
            Expr::Atom(atom) => match atom.inner() {
                Atom::Name(name) => {
                    self.scope_mut().insert_local(name);
                }
                // TODO symbols
                _ => { /* no-op on non-names */ }
            },
        }
    }
}