use crate::protocol::ast::*;
use crate::protocol::input_source::ParseError;
use crate::protocol::parser::{type_table::*, Module};
use crate::protocol::symbol_table::{SymbolTable};

type Unit = ();
pub(crate) type VisitorResult = Result<Unit, ParseError>;

/// Globally configured capacity for large-ish buffers in visitor impls
pub(crate) const BUFFER_INIT_CAP_LARGE: usize = 256;
/// Globally configured capacity for small-ish buffers in visitor impls
pub(crate) const BUFFER_INIT_CAP_SMALL: usize = 64;

/// General context structure that is used while traversing the AST.
pub(crate) struct Ctx<'p> {
    pub heap: &'p mut Heap,
    pub modules: &'p mut [Module],
    pub module_idx: usize, // currently considered module
    pub symbols: &'p mut SymbolTable,
    pub types: &'p mut TypeTable,
    pub arch: &'p crate::protocol::TargetArch,
}

impl<'p> Ctx<'p> {
    /// Returns module `modules[module_idx]`
    pub(crate) fn module(&self) -> &Module {
        &self.modules[self.module_idx]
    }

    pub(crate) fn module_mut(&mut self) -> &mut Module {
        &mut self.modules[self.module_idx]
    }
}

/// Implements the logic that checks the statement union retrieved from the
/// AST and calls the appropriate visit function. This entire macro assumes that
/// `$this` points to `self`, `$stmt` is the statement of type `Statement`,
/// `$ctx` is the context passed to all the visitor calls (of the form
/// `visit_x_stmt(context, id)`) and `$default_return` is the default return
/// value for the statements that will not be visited.
macro_rules! visitor_recursive_statement_impl {
    ($this:expr, $stmt:expr, $ctx:expr, $default_return:expr) => {
        match $stmt {
            Statement::Block(stmt) => {
                let this = stmt.this;
                $this.visit_block_stmt($ctx, this)
            },
            Statement::EndBlock(_stmt) => $default_return,
            Statement::Local(stmt) => {
                let this = stmt.this();
                $this.visit_local_stmt($ctx, this)
            },
            Statement::Labeled(stmt) => {
                let this = stmt.this;
                $this.visit_labeled_stmt($ctx, this)
            },
            Statement::If(stmt) => {
                let this = stmt.this;
                $this.visit_if_stmt($ctx, this)
            },
            Statement::EndIf(_stmt) => $default_return,
            Statement::While(stmt) => {
                let this = stmt.this;
                $this.visit_while_stmt($ctx, this)
            },
            Statement::EndWhile(_stmt) => $default_return,
            Statement::Break(stmt) => {
                let this = stmt.this;
                $this.visit_break_stmt($ctx, this)
            },
            Statement::Continue(stmt) => {
                let this = stmt.this;
                $this.visit_continue_stmt($ctx, this)
            },
            Statement::Synchronous(stmt) => {
                let this = stmt.this;
                $this.visit_synchronous_stmt($ctx, this)
            },
            Statement::EndSynchronous(_stmt) => $default_return,
            Statement::Fork(stmt) => {
                let this = stmt.this;
                $this.visit_fork_stmt($ctx, this)
            },
            Statement::EndFork(_stmt) => $default_return,
            Statement::Select(stmt) => {
                let this = stmt.this;
                $this.visit_select_stmt($ctx, this)
            },
            Statement::EndSelect(_stmt) => $default_return,
            Statement::Return(stmt) => {
                let this = stmt.this;
                $this.visit_return_stmt($ctx, this)
            },
            Statement::Goto(stmt) => {
                let this = stmt.this;
                $this.visit_goto_stmt($ctx, this)
            },
            Statement::New(stmt) => {
                let this = stmt.this;
                $this.visit_new_stmt($ctx, this)
            },
            Statement::Expression(stmt) => {
                let this = stmt.this;
                $this.visit_expr_stmt($ctx, this)
            }
        }
    };
}

macro_rules! visitor_recursive_local_impl {
    ($this:expr, $local:expr, $ctx:expr) => {
        match $local {
            LocalStatement::Channel(local) => {
                let this = local.this;
                $this.visit_local_channel_stmt($ctx, this)
            },
            LocalStatement::Memory(local) => {
                let this = local.this;
                $this.visit_local_memory_stmt($ctx, this)
            }
        }
    }
}

macro_rules! visitor_recursive_definition_impl {
    ($this:expr, $definition:expr, $ctx:expr) => {
        match $definition {
            Definition::Enum(def) => {
                let def = def.this;
                $this.visit_enum_definition($ctx, def)
            },
            Definition::Union(def) => {
                let def = def.this;
                $this.visit_union_definition($ctx, def)
            },
            Definition::Struct(def) => {
                let def = def.this;
                $this.visit_struct_definition($ctx, def)
            },
            Definition::Procedure(def) => {
                let def = def.this;
                $this.visit_procedure_definition($ctx, def)
            },
        }
    }
}

macro_rules! visitor_recursive_expression_impl {
    ($this:expr, $expression:expr, $ctx:expr) => {
        match $expression {
            Expression::Assignment(expr) => {
                let this = expr.this;
                $this.visit_assignment_expr($ctx, this)
            },
            Expression::Binding(expr) => {
                let this = expr.this;
                $this.visit_binding_expr($ctx, this)
            },
            Expression::Conditional(expr) => {
                let this = expr.this;
                $this.visit_conditional_expr($ctx, this)
            },
            Expression::Binary(expr) => {
                let this = expr.this;
                $this.visit_binary_expr($ctx, this)
            },
            Expression::Unary(expr) => {
                let this = expr.this;
                $this.visit_unary_expr($ctx, this)
            },
            Expression::Indexing(expr) => {
                let this = expr.this;
                $this.visit_indexing_expr($ctx, this)
            },
            Expression::Slicing(expr) => {
                let this = expr.this;
                $this.visit_slicing_expr($ctx, this)
            },
            Expression::Select(expr) => {
                let this = expr.this;
                $this.visit_select_expr($ctx, this)
            },
            Expression::Literal(expr) => {
                let this = expr.this;
                $this.visit_literal_expr($ctx, this)
            },
            Expression::Cast(expr) => {
                let this = expr.this;
                $this.visit_cast_expr($ctx, this)
            },
            Expression::Call(expr) => {
                let this = expr.this;
                $this.visit_call_expr($ctx, this)
            },
            Expression::Variable(expr) => {
                let this = expr.this;
                $this.visit_variable_expr($ctx, this)
            },
        }
    };
}

/// Visitor is a generic trait that will fully walk the AST. The default
/// implementation of the visitors is to not recurse. The exception is the
/// top-level `visit_definition`, `visit_stmt` and `visit_expr` methods, which
/// call the appropriate visitor function.
pub(crate) trait Visitor {
    // Entry point
    fn visit_module(&mut self, ctx: &mut Ctx) -> VisitorResult {
        let mut def_index = 0;
        let module_root_id = ctx.modules[ctx.module_idx].root_id;
        loop {
            let definition_id = {
                let root = &ctx.heap[module_root_id];
                if def_index >= root.definitions.len() {
                    return Ok(())
                }

                root.definitions[def_index]
            };

            self.visit_definition(ctx, definition_id)?;
            def_index += 1;
        }
    }

    // Definitions
    // --- enum matching
    fn visit_definition(&mut self, ctx: &mut Ctx, id: DefinitionId) -> VisitorResult {
        return visitor_recursive_definition_impl!(self, &ctx.heap[id], ctx);
    }

    // --- enum variant handling
    fn visit_enum_definition(&mut self, _ctx: &mut Ctx, _id: EnumDefinitionId) -> VisitorResult { Ok(()) }
    fn visit_union_definition(&mut self, _ctx: &mut Ctx, _id: UnionDefinitionId) -> VisitorResult{ Ok(()) }
    fn visit_struct_definition(&mut self, _ctx: &mut Ctx, _id: StructDefinitionId) -> VisitorResult { Ok(()) }
    fn visit_procedure_definition(&mut self, _ctx: &mut Ctx, _id: ProcedureDefinitionId) -> VisitorResult { Ok(()) }

    // Statements
    // --- enum matching
    fn visit_stmt(&mut self, ctx: &mut Ctx, id: StatementId) -> VisitorResult {
        return visitor_recursive_statement_impl!(self, &ctx.heap[id], ctx, Ok(()));
    }

    fn visit_local_stmt(&mut self, ctx: &mut Ctx, id: LocalStatementId) -> VisitorResult {
        return visitor_recursive_local_impl!(self, &ctx.heap[id], ctx);
    }

    // --- enum variant handling
    fn visit_block_stmt(&mut self, _ctx: &mut Ctx, _id: BlockStatementId) -> VisitorResult { Ok(()) }
    fn visit_local_memory_stmt(&mut self, _ctx: &mut Ctx, _id: MemoryStatementId) -> VisitorResult { Ok(()) }
    fn visit_local_channel_stmt(&mut self, _ctx: &mut Ctx, _id: ChannelStatementId) -> VisitorResult { Ok(()) }
    fn visit_labeled_stmt(&mut self, _ctx: &mut Ctx, _id: LabeledStatementId) -> VisitorResult { Ok(()) }
    fn visit_if_stmt(&mut self, _ctx: &mut Ctx, _id: IfStatementId) -> VisitorResult { Ok(()) }
    fn visit_while_stmt(&mut self, _ctx: &mut Ctx, _id: WhileStatementId) -> VisitorResult { Ok(()) }
    fn visit_break_stmt(&mut self, _ctx: &mut Ctx, _id: BreakStatementId) -> VisitorResult { Ok(()) }
    fn visit_continue_stmt(&mut self, _ctx: &mut Ctx, _id: ContinueStatementId) -> VisitorResult { Ok(()) }
    fn visit_synchronous_stmt(&mut self, _ctx: &mut Ctx, _id: SynchronousStatementId) -> VisitorResult { Ok(()) }
    fn visit_fork_stmt(&mut self, _ctx: &mut Ctx, _id: ForkStatementId) -> VisitorResult { Ok(()) }
    fn visit_select_stmt(&mut self, _ctx: &mut Ctx, _id: SelectStatementId) -> VisitorResult { Ok(()) }
    fn visit_return_stmt(&mut self, _ctx: &mut Ctx, _id: ReturnStatementId) -> VisitorResult { Ok(()) }
    fn visit_goto_stmt(&mut self, _ctx: &mut Ctx, _id: GotoStatementId) -> VisitorResult { Ok(()) }
    fn visit_new_stmt(&mut self, _ctx: &mut Ctx, _id: NewStatementId) -> VisitorResult { Ok(()) }
    fn visit_expr_stmt(&mut self, _ctx: &mut Ctx, _id: ExpressionStatementId) -> VisitorResult { Ok(()) }

    // Expressions
    // --- enum matching
    fn visit_expr(&mut self, ctx: &mut Ctx, id: ExpressionId) -> VisitorResult {
        return visitor_recursive_expression_impl!(self, &ctx.heap[id], ctx);
    }

    fn visit_assignment_expr(&mut self, _ctx: &mut Ctx, _id: AssignmentExpressionId) -> VisitorResult { Ok(()) }
    fn visit_binding_expr(&mut self, _ctx: &mut Ctx, _id: BindingExpressionId) -> VisitorResult { Ok(()) }
    fn visit_conditional_expr(&mut self, _ctx: &mut Ctx, _id: ConditionalExpressionId) -> VisitorResult { Ok(()) }
    fn visit_binary_expr(&mut self, _ctx: &mut Ctx, _id: BinaryExpressionId) -> VisitorResult { Ok(()) }
    fn visit_unary_expr(&mut self, _ctx: &mut Ctx, _id: UnaryExpressionId) -> VisitorResult { Ok(()) }
    fn visit_indexing_expr(&mut self, _ctx: &mut Ctx, _id: IndexingExpressionId) -> VisitorResult { Ok(()) }
    fn visit_slicing_expr(&mut self, _ctx: &mut Ctx, _id: SlicingExpressionId) -> VisitorResult { Ok(()) }
    fn visit_select_expr(&mut self, _ctx: &mut Ctx, _id: SelectExpressionId) -> VisitorResult { Ok(()) }
    fn visit_literal_expr(&mut self, _ctx: &mut Ctx, _id: LiteralExpressionId) -> VisitorResult { Ok(()) }
    fn visit_cast_expr(&mut self, _ctx: &mut Ctx, _id: CastExpressionId) -> VisitorResult { Ok(()) }
    fn visit_call_expr(&mut self, _ctx: &mut Ctx, _id: CallExpressionId) -> VisitorResult { Ok(()) }
    fn visit_variable_expr(&mut self, _ctx: &mut Ctx, _id: VariableExpressionId) -> VisitorResult { Ok(()) }
}