//! JavaScript AST structures
//!
//! A [`Program`] is the root of an AST.
//! A `Program` contains a [`BlockStatement`] which is a sequence of [`Statement`]s.
//! A [`Statement`] is a wrapper for [`Stmt`] enum with some additional information.
//! A [`Stmt`] enumerates every possible statement, e.g. [`VariableDeclaration`], [`IfStatement`],
//! [`ExpressionStatement`], etc.
//!
//! [`ExpressionStatement`] is an [`Expression`] wrapped into a statement.
//! An [`Expression`] is a wrapper for [`Expr`] enum that enumerates all possible expressions, e.g.
//! [`CallExpression`], [`MemberExpression`], [`Identifier`], [`AssignmentExpression`], etc.
//!
//! One of [`Expr`] variants is [`FunctionExpression`] that wraps a [`Function`] with its own
//! `BlockStatement` body.

/// [`Expression`], [`Expr`], all expression structs
pub mod expr;

/// [`Statement`], [`Stmt`], all statement structs
pub mod stmt;

/// AST builder DSL
pub mod build;

mod display;

use crate::prelude::*;

use crate::error::{
    ParseError,
    ParseResult,
};

pub use self::expr::*;
pub use self::stmt::*;

/// Represents a complete top-level JS script.
///
/// The only way to create a `Program` externally is to use [`LexicalContext::new_program`],
/// maybe via [`Program::parse_from`] or [`build::script`] helpers. This ensures that
/// its bindings are statically analysed.
#[derive(Debug)]
pub struct Program {
    pub body: BlockStatement, // also contains let/const bindings
    globals: Vec<Binding>,    // everything that goes into `global` scope: variables and functions
}

impl Program {
    pub fn globals_iter(&self) -> impl Iterator<Item = &Binding> {
        self.globals.iter()
    }

    pub fn bindings_iter(&self) -> BlockBindings {
        self.body.bindings_iter()
    }
}

impl PartialEq for Program {
    fn eq(&self, other: &Self) -> bool {
        self.body == other.body
    }
}

impl Eq for Program {}

impl TryFrom<Expression> for Program {
    type Error = ParseError;

    fn try_from(expr: Expression) -> Result<Self, Self::Error> {
        LexicalContext::new_program(|ctx| ctx.block(|_| Ok(vec![Statement::from(expr)])))
    }
}

/// Represents a lexical bindings of a `name`. It might be a hoisted [`Function`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Binding {
    kind: DeclarationKind,
    name: Identifier,
    func: Option<Rc<Function>>,
}

impl Binding {
    pub fn name(&self) -> &Identifier {
        &self.name
    }

    pub fn as_func(&self) -> Option<FunctionExpression> {
        self.func.clone().map(|func| FunctionExpression { func })
    }

    pub fn as_letconst(&self) -> Option<LetConst> {
        use DeclarationKind::*;
        if self.kind != Let && self.kind != Const {
            return None;
        }
        let is_const = self.kind == Const;
        Some(LetConst {
            name: self.name.clone(),
            is_const,
        })
    }

    pub fn is_const(&self) -> bool {
        self.kind == DeclarationKind::Const
    }
}

impl Default for Binding {
    fn default() -> Self {
        Binding {
            kind: DeclarationKind::Var, // to avoid the duplication headache
            name: Identifier::from(""),
            func: None,
        }
    }
}

impl From<FuncBinding> for Binding {
    fn from(fb: FuncBinding) -> Self {
        Binding {
            kind: DeclarationKind::Var,
            name: fb.name,
            func: Some(fb.func),
        }
    }
}

impl From<LetConst> for Binding {
    fn from(lb: LetConst) -> Self {
        let kind = if lb.is_const {
            DeclarationKind::Const
        } else {
            DeclarationKind::Let
        };
        Binding {
            kind,
            name: lb.name,
            func: None,
        }
    }
}

/// A let/const binding.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LetConst {
    name: Identifier,
    is_const: bool,
}

/// A binding to a hoisted [`Function`].
#[derive(Debug, Clone)]
pub struct FuncBinding {
    name: Identifier,
    func: Rc<Function>,
}

impl FuncBinding {
    pub fn name(&self) -> &str {
        self.name.as_str()
    }
}

/// `LexicalContext` collects lexical scope information for some static analysis.
///
/// It produces [`Program`]s, [`Function`]s and [`BlockStatement`]s, ensuring that their bindings
/// and free variables are statically known and correct.
#[derive(Debug)]
pub struct LexicalContext {
    declarations: Vec<Binding>,
    let_names: HashSet<Identifier>, // let/const bindings
    used_ids: HashSet<Identifier>,  // note: they are not free before they leave the scope
}

impl LexicalContext {
    /// Start building a AST for a [`Program`]
    ///
    /// E.g.
    /// ```
    /// use sljs::ast::{self, build::*};
    ///
    /// let program = ast::LexicalContext::new_program(|ctx|
    ///     ctx.enter_block(|ctx| {
    ///         Ok(vec![                                // : Vec<Statement>
    ///             ctx.var_("x", lit(2))?,             // : Statement
    ///             add(ctx.id("x"), lit(3)).into(),    // : Statement
    ///         ])
    ///     })                                          // : ParseResult<BlockStatement>
    /// ).expect("Program");
    /// ```
    pub fn new_program<F>(action: F) -> ParseResult<Program>
    where
        F: FnOnce(&mut LexicalContext) -> ParseResult<BlockStatement>,
    {
        let mut this = LexicalContext::default();
        let mut body = action(&mut this)?;

        let functions = core::mem::take(&mut body.functions);
        let mut globals = this.declarations; // should be vars only
        globals.extend(functions.into_iter().map(Binding::from));

        Ok(Program { body, globals })
    }

    /// AST-builder wrapper for [`Self::new_id`]
    pub fn id(&mut self, name: &str) -> Expression {
        Expression::from(self.new_id(JSString::from(name)))
    }

    /// Create and note the usage of an [`Identifier`] in an expression context.
    pub fn new_id(&mut self, name: JSString) -> Identifier {
        let id = Identifier(name);
        self.used_ids.insert(id.clone());
        id
    }

    /// AST-builder wrapper for [`Self::declare_var`] for one variable without initialization
    pub fn var(&mut self, name: &str) -> ParseResult<Statement> {
        let kind = DeclarationKind::Var;
        let name = Identifier::from(name);
        self.declare_var(kind, &name)?;
        let declarations = vec![(name, None)];
        Ok(Statement::from(VariableDeclaration { kind, declarations }))
    }

    /// AST-builder wrapper for [`Self::declare_var`] for one variable with initialization
    pub fn var_(&mut self, name: &str, initexpr: Expression) -> ParseResult<Statement> {
        let kind = DeclarationKind::Var;
        let name = Identifier::from(name);
        self.declare_var(kind, &name)?;
        let declarations = vec![(name, Some(initexpr))];
        Ok(Statement::from(VariableDeclaration { kind, declarations }))
    }

    /// Note a let/const/var declaration.
    pub fn declare_var(&mut self, kind: DeclarationKind, name: &Identifier) -> ParseResult<()> {
        let binding = Binding {
            kind,
            name: name.clone(),
            func: None,
        };
        if let Some(lb) = binding.as_letconst() {
            if !self.let_names.insert(lb.name) {
                return Err(ParseError::BindingRedeclared { name: name.clone() });
            }
        }
        self.declarations.push(binding);
        Ok(())
    }

    /// AST-builder wrapper for [`Self::enter_function`]
    ///
    /// It takes a closure that produces the list of statements in the body
    /// and return a [`FunctionExpression`] wrapped in an [`Expression`]:
    /// ```
    /// use sljs::{
    ///     ast::LexicalContext,
    ///     ast::build::*,
    /// };
    ///
    /// let mut ctx = LexicalContext::default();
    /// let sqr = ctx.function(&["x"], |ctx| Ok(vec![
    ///     return_(mul(ctx.id("x"), ctx.id("x"))),
    /// ])).expect("FunctionExpression");
    /// ```
    pub fn function<F>(&mut self, argnames: &[&str], body: F) -> ParseResult<Expression>
    where
        F: FnMut(&mut LexicalContext) -> ParseResult<Vec<Statement>>,
    {
        let func = self.enter_function(|ctx| {
            let mut params = vec![];
            for &arg in argnames {
                ctx.var(arg)?;
                params.push(Identifier::from(arg));
            }
            let body = ctx.enter_block(body)?;
            Ok((None, params, body))
        })?;
        let func = Rc::new(func);
        Ok(Expression::from(FunctionExpression { func }))
    }

    /// Note a function declaration.
    pub fn declare_func(&mut self, function: &FunctionDeclaration) -> ParseResult<()> {
        self.declarations.push(Binding {
            kind: DeclarationKind::Var,
            name: function.get_name(),
            func: Some(Rc::clone(&function.func)),
        });
        Ok(())
    }

    /// An AST-builder wrapper for [`Self::enter_block`]
    pub fn block<F>(&mut self, action: F) -> ParseResult<BlockStatement>
    where
        F: FnOnce(&mut Self) -> ParseResult<Vec<Statement>>,
    {
        self.enter_block(action)
    }

    pub fn enter_block<F>(&mut self, action: F) -> ParseResult<BlockStatement>
    where
        F: FnOnce(&mut LexicalContext) -> ParseResult<Vec<Statement>>,
    {
        // inner_ctx accumulates used identifiers and declared bindings.
        let mut inner_ctx = LexicalContext::default();

        let result = action(&mut inner_ctx);

        let (functions, bindings) = self.consume(inner_ctx)?;
        let body = result?;
        Ok(BlockStatement {
            body,
            functions,
            bindings,
        })
    }

    pub fn enter_function<F>(&mut self, action: F) -> ParseResult<Function>
    where
        F: FnOnce(
            &mut LexicalContext,
        ) -> ParseResult<(Option<Identifier>, Vec<Pattern>, BlockStatement)>,
    {
        let mut ctx = LexicalContext::default();
        ctx.var("arguments")?;

        let (id, params, body) = action(&mut ctx)?;

        // body.functions contain functions
        // body.bindings contain let/const
        // ctx.declarations contain vars
        // ctx.used_ids

        let locals = ctx.drain_declarations()?;
        let free_variables = Vec::from_iter(ctx.used_ids.iter().cloned());

        Ok(Function {
            id,
            params,
            body,
            locals,
            free_variables,
        })
    }

    /// Split bindings: leave vars in `self`, split off everything else.
    ///
    /// The result:
    /// - starts with all function bindings in undefined order
    /// - const/lets come later in the order of their declaration
    fn split_off_block_bindings(&mut self) -> ParseResult<(Vec<FuncBinding>, Vec<LetConst>)> {
        // remove split bindings from used_ids
        // the order and names of vars does not matter
        // function binding names cannot overlap with const/let bindings.
        // the order of function bindings relative to each other does not matter
        // function binding names can repeat, only the last binding is used.
        // the order of let/const-bindings relative to each other must be preserved
        // const/let binding names cannot repeat.
        use DeclarationKind::*;

        let mut lets = Vec::new();
        let mut funcs = HashMap::<Identifier, Rc<Function>>::new();
        let mut var_at = 0;

        for at in 0..self.declarations.len() {
            let name = self.declarations[at].name.clone();
            self.used_ids.remove(&name);

            if let Some(func) = self.declarations[at].func.take() {
                if self.let_names.contains(&name) {
                    return Err(ParseError::BindingRedeclared { name });
                }
                funcs.insert(name, func);
                continue;
            }
            let kind = self.declarations[at].kind;
            match kind {
                Var => {
                    if var_at != at {
                        self.declarations[var_at] = Binding::default();
                        self.declarations.as_mut_slice().swap(var_at, at);
                    }
                    var_at += 1;
                }
                Let | Const => {
                    // name uniqueness is checked in Self::declare_var()
                    let is_const = kind == Const;
                    lets.push(LetConst { name, is_const })
                }
            }
        }

        let funcs = Vec::from_iter(funcs.drain().map(|(name, func)| FuncBinding { name, func }));

        self.declarations.truncate(var_at);
        self.let_names = HashSet::new(); // drop old names

        Ok((funcs, lets))
    }

    fn consume(&mut self, mut other: Self) -> ParseResult<(Vec<FuncBinding>, Vec<LetConst>)> {
        // extract all function and let/const bindings
        let (functions, bindings) = other.split_off_block_bindings()?;

        // add all remaining usages to the outer used_variables
        self.used_ids.extend(other.used_ids.into_iter());
        self.declarations.extend(other.declarations.into_iter()); // vars
        Ok((functions, bindings))
    }

    fn drain_declarations(&mut self) -> ParseResult<Vec<Binding>> {
        let vars = core::mem::take(&mut self.declarations);
        for b in vars.iter() {
            self.used_ids.remove(b.name());
        }
        Ok(vars)
    }
}

impl Default for LexicalContext {
    /// This is not what you usually want, build a [`Program`] with [`LexicalContext::new_program`]
    /// instead.
    fn default() -> Self {
        LexicalContext {
            declarations: Vec::new(),
            let_names: HashSet::new(),
            used_ids: HashSet::new(),
        }
    }
}