#[macro_use] mod visitor;
pub(crate) mod symbol_table;
pub(crate) mod type_table;
pub(crate) mod tokens;
pub(crate) mod token_parsing;
pub(crate) mod pass_tokenizer;
pub(crate) mod pass_symbols;
pub(crate) mod pass_imports;
pub(crate) mod pass_definitions;
pub(crate) mod pass_definitions_types;
pub(crate) mod pass_validation_linking;
pub(crate) mod pass_rewriting;
pub(crate) mod pass_typing;
pub(crate) mod pass_stack_size;

use tokens::*;
use crate::collections::*;
use visitor::Visitor;
use pass_tokenizer::PassTokenizer;
use pass_symbols::PassSymbols;
use pass_imports::PassImport;
use pass_definitions::PassDefinitions;
use pass_validation_linking::PassValidationLinking;
use pass_typing::{PassTyping, ResolveQueue};
use pass_rewriting::PassRewriting;
use pass_stack_size::PassStackSize;
use symbol_table::*;
use type_table::*;

use crate::protocol::ast::*;
use crate::protocol::input_source::*;

use crate::protocol::ast_printer::ASTWriter;
use crate::protocol::parser::type_table::PolymorphicVariable;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum ModuleCompilationPhase {
    Tokenized,              // source is tokenized
    SymbolsScanned,         // all definitions are linked to their type class
    ImportsResolved,        // all imports are added to the symbol table
    DefinitionsParsed,      // produced the AST for the entire module
    TypesAddedToTable,      // added all definitions to the type table
    ValidatedAndLinked,     // AST is traversed and has linked the required AST nodes
    Typed,                  // Type inference and checking has been performed
    Rewritten,              // Special AST nodes are rewritten into regular AST nodes
    // When we continue with the compiler:
    // StackSize
}

pub struct Module {
    // Buffers
    pub source: InputSource,
    pub tokens: TokenBuffer,
    // Identifiers
    pub root_id: RootId,
    pub name: Option<(PragmaId, StringRef<'static>)>,
    pub version: Option<(PragmaId, i64)>,
    pub phase: ModuleCompilationPhase,
}

pub struct TargetArch {
    pub void_type_id: TypeId,
    pub message_type_id: TypeId,
    pub bool_type_id: TypeId,
    pub uint8_type_id: TypeId,
    pub uint16_type_id: TypeId,
    pub uint32_type_id: TypeId,
    pub uint64_type_id: TypeId,
    pub sint8_type_id: TypeId,
    pub sint16_type_id: TypeId,
    pub sint32_type_id: TypeId,
    pub sint64_type_id: TypeId,
    pub char_type_id: TypeId,
    pub string_type_id: TypeId,
    pub array_type_id: TypeId,
    pub slice_type_id: TypeId,
    pub input_type_id: TypeId,
    pub output_type_id: TypeId,
    pub pointer_type_id: TypeId,
}

impl TargetArch {
    fn new() -> Self {
        return Self{
            void_type_id: TypeId::new_invalid(),
            bool_type_id: TypeId::new_invalid(),
            message_type_id: TypeId::new_invalid(),
            uint8_type_id: TypeId::new_invalid(),
            uint16_type_id: TypeId::new_invalid(),
            uint32_type_id: TypeId::new_invalid(),
            uint64_type_id: TypeId::new_invalid(),
            sint8_type_id: TypeId::new_invalid(),
            sint16_type_id: TypeId::new_invalid(),
            sint32_type_id: TypeId::new_invalid(),
            sint64_type_id: TypeId::new_invalid(),
            char_type_id: TypeId::new_invalid(),
            string_type_id: TypeId::new_invalid(),
            array_type_id: TypeId::new_invalid(),
            slice_type_id: TypeId::new_invalid(),
            input_type_id: TypeId::new_invalid(),
            output_type_id: TypeId::new_invalid(),
            pointer_type_id: TypeId::new_invalid(),
        }
    }
}

pub struct PassCtx<'a> {
    heap: &'a mut Heap,
    symbols: &'a mut SymbolTable,
    pool: &'a mut StringPool,
    arch: &'a TargetArch,
}

pub struct Parser {
    // Storage of all information created/gathered during compilation.
    pub(crate) heap: Heap,
    pub(crate) string_pool: StringPool, // Do not deallocate, holds all strings
    pub(crate) modules: Vec<Module>,
    pub(crate) symbol_table: SymbolTable,
    pub(crate) type_table: TypeTable,
    // Compiler passes, used as little state machine that keep their memory
    // around.
    pass_tokenizer: PassTokenizer,
    pass_symbols: PassSymbols,
    pass_import: PassImport,
    pass_definitions: PassDefinitions,
    pass_validation: PassValidationLinking,
    pass_typing: PassTyping,
    pass_rewriting: PassRewriting,
    pass_stack_size: PassStackSize,
    // Compiler options
    pub write_ast_to: Option<String>,
    pub(crate) arch: TargetArch,
}

impl Parser {
    pub fn new() -> Self {
        let mut parser = Parser{
            heap: Heap::new(),
            string_pool: StringPool::new(),
            modules: Vec::new(),
            symbol_table: SymbolTable::new(),
            type_table: TypeTable::new(),
            pass_tokenizer: PassTokenizer::new(),
            pass_symbols: PassSymbols::new(),
            pass_import: PassImport::new(),
            pass_definitions: PassDefinitions::new(),
            pass_validation: PassValidationLinking::new(),
            pass_typing: PassTyping::new(),
            pass_rewriting: PassRewriting::new(),
            pass_stack_size: PassStackSize::new(),
            write_ast_to: None,
            arch: TargetArch::new(),
        };

        parser.symbol_table.insert_scope(None, SymbolScope::Global);

        // Insert builtin types
        // TODO: At some point use correct values for size/alignment
        parser.arch.void_type_id    = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Void], false, 0, 1);
        parser.arch.message_type_id = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Message], false, 24, 8);
        parser.arch.bool_type_id    = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Bool], false, 1, 1);
        parser.arch.uint8_type_id   = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::UInt8], false, 1, 1);
        parser.arch.uint16_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::UInt16], false, 2, 2);
        parser.arch.uint32_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::UInt32], false, 4, 4);
        parser.arch.uint64_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::UInt64], false, 8, 8);
        parser.arch.sint8_type_id   = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::SInt8], false, 1, 1);
        parser.arch.sint16_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::SInt16], false, 2, 2);
        parser.arch.sint32_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::SInt32], false, 4, 4);
        parser.arch.sint64_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::SInt64], false, 8, 8);
        parser.arch.char_type_id    = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Character], false, 4, 4);
        parser.arch.string_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::String], false, 24, 8);
        parser.arch.array_type_id   = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Array, ConcreteTypePart::Void], true, 24, 8);
        parser.arch.slice_type_id   = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Slice, ConcreteTypePart::Void], true, 16, 4);
        parser.arch.input_type_id   = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Input, ConcreteTypePart::Void], true, 8, 8);
        parser.arch.output_type_id  = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Output, ConcreteTypePart::Void], true, 8, 8);
        parser.arch.pointer_type_id = insert_builtin_type(&mut parser.type_table, vec![ConcreteTypePart::Pointer, ConcreteTypePart::Void], true, 8, 8);

        // Insert builtin functions
        fn quick_type(variants: &[ParserTypeVariant]) -> ParserType {
            let mut t = ParserType{ elements: Vec::with_capacity(variants.len()), full_span: InputSpan::new() };
            for variant in variants {
                t.elements.push(ParserTypeElement{ element_span: InputSpan::new(), variant: variant.clone() });
            }
            t
        }

        use ParserTypeVariant as PTV;
        insert_builtin_function(&mut parser, "get", &["T"], |id| (
            vec![
                ("input", quick_type(&[PTV::Input, PTV::PolymorphicArgument(id.upcast(), 0)]))
            ],
            quick_type(&[PTV::PolymorphicArgument(id.upcast(), 0)])
        ));
        insert_builtin_function(&mut parser, "put", &["T"], |id| (
            vec![
                ("output", quick_type(&[PTV::Output, PTV::PolymorphicArgument(id.upcast(), 0)])),
                ("value", quick_type(&[PTV::PolymorphicArgument(id.upcast(), 0)])),
            ],
            quick_type(&[PTV::Void])
        ));
        insert_builtin_function(&mut parser, "fires", &["T"], |id| (
            vec![
                ("port", quick_type(&[PTV::InputOrOutput, PTV::PolymorphicArgument(id.upcast(), 0)]))
            ],
            quick_type(&[PTV::Bool])
        ));
        insert_builtin_function(&mut parser, "create", &["T"], |id| (
            vec![
                ("length", quick_type(&[PTV::IntegerLike]))
            ],
            quick_type(&[PTV::ArrayLike, PTV::PolymorphicArgument(id.upcast(), 0)])
        ));
        insert_builtin_function(&mut parser, "length", &["T"], |id| (
            vec![
                ("array", quick_type(&[PTV::ArrayLike, PTV::PolymorphicArgument(id.upcast(), 0)]))
            ],
            quick_type(&[PTV::UInt32]) // TODO: @PtrInt
        ));
        insert_builtin_function(&mut parser, "assert", &[], |_id| (
            vec![
                ("condition", quick_type(&[PTV::Bool])),
            ],
            quick_type(&[PTV::Void])
        ));
        insert_builtin_function(&mut parser, "print", &[], |_id| (
            vec![
                ("message", quick_type(&[PTV::String])),
            ],
            quick_type(&[PTV::Void])
        ));

        parser
    }

    pub fn feed(&mut self, mut source: InputSource) -> Result<(), ParseError> {
        let mut token_buffer = TokenBuffer::new();
        self.pass_tokenizer.tokenize(&mut source, &mut token_buffer)?;

        let module = Module{
            source,
            tokens: token_buffer,
            root_id: RootId::new_invalid(),
            name: None,
            version: None,
            phase: ModuleCompilationPhase::Tokenized,
        };
        self.modules.push(module);

        Ok(())
    }

    pub fn parse(&mut self) -> Result<(), ParseError> {
        let mut pass_ctx = PassCtx{
            heap: &mut self.heap,
            symbols: &mut self.symbol_table,
            pool: &mut self.string_pool,
            arch: &self.arch,
        };

        // Advance all modules to the phase where all symbols are scanned
        for module_idx in 0..self.modules.len() {
            self.pass_symbols.parse(&mut self.modules, module_idx, &mut pass_ctx)?;
        }

        // With all symbols scanned, perform further compilation until we can
        // add all base types to the type table.
        for module_idx in 0..self.modules.len() {
            self.pass_import.parse(&mut self.modules, module_idx, &mut pass_ctx)?;
            self.pass_definitions.parse(&mut self.modules, module_idx, &mut pass_ctx)?;
        }

        // Add every known type to the type table
        self.type_table.build_base_types(&mut self.modules, &mut pass_ctx)?;

        // Continue compilation with the remaining phases now that the types
        // are all in the type table
        for module_idx in 0..self.modules.len() {
            let mut ctx = visitor::Ctx{
                heap: &mut self.heap,
                modules: &mut self.modules,
                module_idx,
                symbols: &mut self.symbol_table,
                types: &mut self.type_table,
                arch: &self.arch,
            };
            self.pass_validation.visit_module(&mut ctx)?;
        }

        // Perform typechecking on all modules
        let mut queue = ResolveQueue::new();
        for module_idx in 0..self.modules.len() {
            let mut ctx = visitor::Ctx{
                heap: &mut self.heap,
                modules: &mut self.modules,
                module_idx,
                symbols: &mut self.symbol_table,
                types: &mut self.type_table,
                arch: &self.arch,
            };
            self.pass_typing.queue_module_definitions(&mut ctx, &mut queue);
        };
        while !queue.is_empty() {
            let top = queue.pop_front().unwrap();
            let mut ctx = visitor::Ctx{
                heap: &mut self.heap,
                modules: &mut self.modules,
                module_idx: top.root_id.index as usize,
                symbols: &mut self.symbol_table,
                types: &mut self.type_table,
                arch: &self.arch,
            };
            self.pass_typing.handle_module_definition(&mut ctx, &mut queue, top)?;
        }

        // Rewrite nodes in tree, then prepare for execution of code
        for module_idx in 0..self.modules.len() {
            self.modules[module_idx].phase = ModuleCompilationPhase::Typed;
            let mut ctx = visitor::Ctx{
                heap: &mut self.heap,
                modules: &mut self.modules,
                module_idx,
                symbols: &mut self.symbol_table,
                types: &mut self.type_table,
                arch: &self.arch,
            };
            self.pass_rewriting.visit_module(&mut ctx)?;
            self.pass_stack_size.visit_module(&mut ctx)?;
        }

        // Write out desired information
        if let Some(filename) = &self.write_ast_to {
            let mut writer = ASTWriter::new();
            let mut file = std::fs::File::create(std::path::Path::new(filename)).unwrap();
            writer.write_ast(&mut file, &self.heap);
        }

        Ok(())
    }
}

fn insert_builtin_type(type_table: &mut TypeTable, parts: Vec<ConcreteTypePart>, has_poly_var: bool, size: usize, alignment: usize) -> TypeId {
    const POLY_VARS: [PolymorphicVariable; 1] = [PolymorphicVariable{
        identifier: Identifier::new_empty(InputSpan::new()),
        is_in_use: false,
    }];

    let concrete_type = ConcreteType{ parts };
    let poly_var = if has_poly_var {
        POLY_VARS.as_slice()
    } else {
        &[]
    };

    return type_table.add_builtin_data_type(concrete_type, poly_var, size, alignment);
}

// Note: args and return type need to be a function because we need to know the function ID.
fn insert_builtin_function<T: Fn(ProcedureDefinitionId) -> (Vec<(&'static str, ParserType)>, ParserType)> (
    p: &mut Parser, func_name: &str, polymorphic: &[&str], arg_and_return_fn: T
) {
    // Insert into AST (to get an ID), also prepare the polymorphic variables
    // we need later for the type table
    let mut ast_poly_vars = Vec::with_capacity(polymorphic.len());
    let mut type_poly_vars = Vec::with_capacity(polymorphic.len());
    for poly_var in polymorphic {
        let identifier = Identifier{ span: InputSpan::new(), value: p.string_pool.intern(poly_var.as_bytes()) } ;
        ast_poly_vars.push(identifier.clone());
        type_poly_vars.push(PolymorphicVariable{ identifier, is_in_use: false });
    }

    let func_ident_ref = p.string_pool.intern(func_name.as_bytes());
    let procedure_id = p.heap.alloc_procedure_definition(|this| ProcedureDefinition {
        this,
        defined_in: RootId::new_invalid(),
        builtin: true,
        kind: ProcedureKind::Function,
        span: InputSpan::new(),
        identifier: Identifier{ span: InputSpan::new(), value: func_ident_ref.clone() },
        poly_vars: ast_poly_vars,
        return_type: None,
        parameters: Vec::new(),
        scope: ScopeId::new_invalid(),
        body: BlockStatementId::new_invalid(),
        monomorphs: Vec::new(),
    });

    // Modify AST with more information about the procedure
    let (arguments, return_type) = arg_and_return_fn(procedure_id);

    let mut parameters = Vec::with_capacity(arguments.len());
    for (arg_name, arg_type) in arguments {
        let identifier = Identifier{ span: InputSpan::new(), value: p.string_pool.intern(arg_name.as_bytes()) };
        let param_id = p.heap.alloc_variable(|this| Variable{
            this,
            kind: VariableKind::Parameter,
            parser_type: arg_type.clone(),
            identifier,
            relative_pos_in_parent: 0,
            unique_id_in_scope: 0
        });
        parameters.push(param_id);
    }

    let func = &mut p.heap[procedure_id];
    func.parameters = parameters;
    func.return_type = Some(return_type);

    // Insert into symbol table
    p.symbol_table.insert_symbol(SymbolScope::Global, Symbol{
        name: func_ident_ref,
        variant: SymbolVariant::Definition(SymbolDefinition{
            defined_in_module: RootId::new_invalid(),
            defined_in_scope: SymbolScope::Global,
            definition_span: InputSpan::new(),
            identifier_span: InputSpan::new(),
            imported_at: None,
            class: DefinitionClass::Function,
            definition_id: procedure_id.upcast(),
        })
    }).unwrap();

    // Insert into type table
    // let mut concrete_type = ConcreteType::default();
    // concrete_type.parts.push(ConcreteTypePart::Function(procedure_id, type_poly_vars.len() as u32));
    //
    // for _ in 0..type_poly_vars.len() {
    //     concrete_type.parts.push(ConcreteTypePart::Void); // doesn't matter (I hope...)
    // }
    // p.type_table.add_builtin_procedure_type(concrete_type, &type_poly_vars);
}