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

use super::symbol_table::*;

use crate::protocol::input_source::{ParseError, InputSpan};

use super::tokens::*;

use super::token_parsing::*;

use super::{Module, ModuleCompilationPhase, PassCtx};

/// Scans the module and finds all module-level type definitions. These will be

/// added to the symbol table such that during AST-construction we know which

/// identifiers point to types. Will also parse all pragmas to determine module

/// names.

pub(crate) struct PassSymbols {

    symbols: Vec<Symbol>,

    pragmas: Vec<PragmaId>,

    imports: Vec<ImportId>,

    definitions: Vec<DefinitionId>,

    buffer: String,

    has_pragma_version: bool,

    has_pragma_module: bool,

}

impl PassSymbols {

    pub(crate) fn new() -> Self {

        Self{

            symbols: Vec::with_capacity(128),

            pragmas: Vec::with_capacity(8),

            imports: Vec::with_capacity(32),

            definitions: Vec::with_capacity(128),

            buffer: String::with_capacity(128),

            has_pragma_version: false,

            has_pragma_module: false,

        }

    }

    fn reset(&mut self) {

        self.symbols.clear();

        self.pragmas.clear();

        self.imports.clear();

        self.definitions.clear();

        self.has_pragma_version = false;

        self.has_pragma_module = false;

    }

    pub(crate) fn parse(&mut self, modules: &mut [Module], module_idx: usize, ctx: &mut PassCtx) -> Result<(), ParseError> {

        self.reset();

        let module = &mut modules[module_idx];

        let module_range = &module.tokens.ranges[0];

        debug_assert_eq!(module.phase, ModuleCompilationPhase::Tokenized);

        debug_assert_eq!(module_range.range_kind, TokenRangeKind::Module);

        debug_assert!(module.root_id.is_invalid()); // not set yet,

        // Preallocate root in the heap

        let root_id = ctx.heap.alloc_protocol_description(|this| {

            Root{

                this,

                pragmas: Vec::new(),

                imports: Vec::new(),

                definitions: Vec::new(),

            }

        });

        module.root_id = root_id;

        // Retrieve first range index, then make immutable borrow

        let mut range_idx = module_range.first_child_idx;

        // Visit token ranges to detect definitions and pragmas

        loop {

            let range_idx_usize = range_idx as usize;

            let cur_range = &module.tokens.ranges[range_idx_usize];

            let next_sibling_idx = cur_range.next_sibling_idx;

            let range_kind = cur_range.range_kind;

            // Parse if it is a definition or a pragma

            if range_kind == TokenRangeKind::Definition {

                self.visit_definition_range(modules, module_idx, ctx, range_idx_usize)?;

            } else if range_kind == TokenRangeKind::Pragma {

                self.visit_pragma_range(modules, module_idx, ctx, range_idx_usize)?;

            }

            if next_sibling_idx == NO_SIBLING {

                break;

            } else {

                range_idx = next_sibling_idx;

            }

        }

        // Add the module's symbol scope and the symbols we just parsed

        let module_scope = SymbolScope::Module(root_id);

        ctx.symbols.insert_scope(Some(SymbolScope::Global), module_scope);

        for symbol in self.symbols.drain(..) {

            ctx.symbols.insert_scope(Some(module_scope), SymbolScope::Definition(symbol.variant.as_definition().definition_id));

            if let Err((new_symbol, old_symbol)) = ctx.symbols.insert_symbol(module_scope, symbol) {

                return Err(construct_symbol_conflict_error(modules, module_idx, ctx, &new_symbol, &old_symbol))

            }

        }

        // Modify the preallocated root

        let root = &mut ctx.heap[root_id];

        root.pragmas.extend(self.pragmas.drain(..));

        root.definitions.extend(self.definitions.drain(..));

        let module = &mut modules[module_idx];

        module.phase = ModuleCompilationPhase::SymbolsScanned;

        Ok(())

    }

        let range = &module.tokens.ranges[range_idx];

        let mut iter = module.tokens.iter_range(range);

        // Consume pragma name

        let (pragma_section, pragma_start, _) = consume_pragma(&module.source, &mut iter)?;

        // Consume pragma values

        if pragma_section == b"#module" {

            // Check if name is defined twice within the same file

            if self.has_pragma_module {

                return Err(ParseError::new_error_str_at_pos(&module.source, pragma_start, "module name is defined twice"));

            }

            // Consume the domain-name

            let (module_name, module_span) = consume_domain_ident(&module.source, &mut iter)?;

            if iter.next().is_some() {

                return Err(ParseError::new_error_str_at_pos(&module.source, iter.last_valid_pos(), "expected end of #module pragma after module name"));

            }

            // Add to heap and symbol table

            let pragma_span = InputSpan::from_positions(pragma_start, module_span.end);

            let module_name = ctx.pool.intern(module_name);

            let pragma_id = ctx.heap.alloc_pragma(|this| Pragma::Module(PragmaModule{

                this,

                span: pragma_span,

                value: Identifier{ span: module_span, value: module_name.clone() },

            }));

            self.pragmas.push(pragma_id);

                // Naming conflict

                let this_module = &modules[module_idx];

                let other_module = seek_module(modules, other_module_root_id).unwrap();

                let other_module_pragma_id = other_module.name.as_ref().map(|v| (*v).0).unwrap();

                let other_pragma = ctx.heap[other_module_pragma_id].as_module();

                return Err(ParseError::new_error_str_at_span(

                    &this_module.source, pragma_span, "conflict in module name"

                ).with_info_str_at_span(

                    &other_module.source, other_pragma.span, "other module is defined here"

                ));

            }

            self.has_pragma_module = true;

        } else if pragma_section == b"#version" {

            // Check if version is defined twice within the same file

            if self.has_pragma_version {

                return Err(ParseError::new_error_str_at_pos(&module.source, pragma_start, "module version is defined twice"));

            }

            // Consume the version pragma

            let (version, version_span) = consume_integer_literal(&module.source, &mut iter, &mut self.buffer)?;

            let pragma_id = ctx.heap.alloc_pragma(|this| Pragma::Version(PragmaVersion{

                this,

                span: InputSpan::from_positions(pragma_start, version_span.end),

                version,

            }));

            self.pragmas.push(pragma_id);

            self.has_pragma_version = true;

        } else {

            // Custom pragma, maybe we support this in the future, but for now

            // we don't.

            return Err(ParseError::new_error_str_at_pos(&module.source, pragma_start, "illegal pragma name"));

        }

        Ok(())

    }

    fn visit_definition_range(&mut self, modules: &[Module], module_idx: usize, ctx: &mut PassCtx, range_idx: usize) -> Result<(), ParseError> {

        let module = &modules[module_idx];

        let range = &module.tokens.ranges[range_idx];

        let definition_span = InputSpan::from_positions(

            module.tokens.start_pos(range),

            module.tokens.end_pos(range)

        );

        let mut iter = module.tokens.iter_range(range);

        // First ident must be type of symbol

        let (kw_text, _) = consume_any_ident(&module.source, &mut iter).unwrap();

        // Retrieve identifier of definition

        let identifier = consume_ident_interned(&module.source, &mut iter, ctx)?;

        let mut poly_vars = Vec::new();

        maybe_consume_comma_separated(

            TokenKind::OpenAngle, TokenKind::CloseAngle, &module.source, &mut iter, ctx,

            |source, iter, ctx| consume_ident_interned(source, iter, ctx),

            &mut poly_vars, "a polymorphic variable", None

        )?;

        let ident_text = identifier.value.clone(); // because we need it later

        let ident_span = identifier.span.clone();

        // Reserve space in AST for definition and add it to the symbol table

        let definition_class;

        let ast_definition_id;

        match kw_text {

            KW_STRUCT => {

                let struct_def_id = ctx.heap.alloc_struct_definition(|this| {

                    StructDefinition::new_empty(this, module.root_id, definition_span, identifier, poly_vars)

                });

                definition_class = DefinitionClass::Struct;

                ast_definition_id = struct_def_id.upcast();

            },

            KW_ENUM => {

                let enum_def_id = ctx.heap.alloc_enum_definition(|this| {

                    EnumDefinition::new_empty(this, module.root_id, definition_span, identifier, poly_vars)

                });

                definition_class = DefinitionClass::Enum;

                ast_definition_id = enum_def_id.upcast();

            },

            KW_UNION => {

                let union_def_id = ctx.heap.alloc_union_definition(|this| {

                    UnionDefinition::new_empty(this, module.root_id, definition_span, identifier, poly_vars)

                });

                definition_class = DefinitionClass::Union;

                ast_definition_id = union_def_id.upcast()

            },

            KW_FUNCTION => {

                let func_def_id = ctx.heap.alloc_function_definition(|this| {

                    FunctionDefinition::new_empty(this, module.root_id, definition_span, identifier, poly_vars)

                });

                definition_class = DefinitionClass::Function;

                ast_definition_id = func_def_id.upcast();

            },

            KW_PRIMITIVE | KW_COMPOSITE => {

                let component_variant = if kw_text == KW_PRIMITIVE {

                    ComponentVariant::Primitive

                } else {

                    ComponentVariant::Composite

                };

                let comp_def_id = ctx.heap.alloc_component_definition(|this| {

                    ComponentDefinition::new_empty(this, module.root_id, definition_span, component_variant, identifier, poly_vars)

                });

                definition_class = DefinitionClass::Component;

                ast_definition_id = comp_def_id.upcast();

            },

            _ => unreachable!("encountered keyword '{}' in definition range", String::from_utf8_lossy(kw_text)),

        }

        let symbol = Symbol{

            name: ident_text,

            variant: SymbolVariant::Definition(SymbolDefinition{

                defined_in_module: module.root_id,

                defined_in_scope: SymbolScope::Module(module.root_id),

                definition_span,

                identifier_span: ident_span,

                imported_at: None,

                class: definition_class,

                definition_id: ast_definition_id,

            }),

        };

        self.symbols.push(symbol);

        self.definitions.push(ast_definition_id);

        Ok(())

    }

}

use std::sync::{Arc, Mutex, Condvar};

use std::sync::atomic::{AtomicU32, AtomicBool, Ordering};

use std::collections::VecDeque;

use crate::protocol::*;

use super::communication::Message;

use super::component::{wake_up_if_sleeping, CompPDL, CompCtx};

use super::store::{ComponentStore, ComponentReservation, QueueDynMpsc, QueueDynProducer};

use super::scheduler::*;

// -----------------------------------------------------------------------------

// Component

// -----------------------------------------------------------------------------

/// Key to a component. Type system somewhat ensures that there can only be one

/// of these. Only with a key one may retrieve privately-accessible memory for

/// a component. Practically just a generational index, like `CompId` is.

#[derive(Debug, Copy, Clone, PartialEq, Eq)]

pub(crate) struct CompKey(pub u32);

impl CompKey {

    pub(crate) fn downgrade(&self) -> CompId {

        return CompId(self.0);

    }

}

/// Generational ID of a component

#[derive(Debug, Copy, Clone, PartialEq, Eq)]

pub struct CompId(pub u32);

impl CompId {

    pub(crate) fn new_invalid() -> CompId {

        return CompId(u32::MAX);

    }

    /// Upgrade component ID to component key. Unsafe because the caller needs

    /// to make sure that only one component key can exist at a time (to ensure

    /// a component can only be scheduled/executed by one thread).

    pub(crate) unsafe fn upgrade(&self) -> CompKey {

        return CompKey(self.0);

    }

}

/// Handle to a component that is being created.

pub(crate) struct CompReserved {

    reservation: ComponentReservation,

}

impl CompReserved {

    pub(crate) fn id(&self) -> CompId {

        return CompId(self.reservation.index)

    }

}

/// Private fields of a component, may only be modified by a single thread at

/// a time.

pub(crate) struct RuntimeComp {

    pub public: CompPublic,

    pub code: CompPDL,

    pub ctx: CompCtx,

    pub inbox: QueueDynMpsc<Message>,

    pub exiting: bool,

}

/// Should contain everything that is accessible in a thread-safe manner

// TODO: Do something about the `num_handles` thing. This needs to be a bit more

//  "foolproof" to lighten the mental burden of using the `num_handles`

//  variable.

pub(crate) struct CompPublic {

    pub sleeping: AtomicBool,

    pub num_handles: AtomicU32, // manually modified (!)

    inbox: QueueDynProducer<Message>,

}

/// Handle to public part of a component. Would be nice if we could

/// automagically manage the `num_handles` counter. But when it reaches zero we

/// need to manually remove the handle from the runtime. So we just have debug

/// code to make sure this actually happens.

pub(crate) struct CompHandle {

    target: *const CompPublic,

    id: CompId, // TODO: @Remove after debugging

    #[cfg(debug_assertions)] decremented: bool,

}

impl CompHandle {

    fn new(id: CompId, public: &CompPublic) -> CompHandle {

        let handle = CompHandle{

            target: public,

            id,

            #[cfg(debug_assertions)] decremented: false,

        };

        handle.increment_users();

        return handle;

    }

    pub(crate) fn send_message(&self, sched_ctx: &SchedulerCtx, message: Message, try_wake_up: bool) {

        sched_ctx.log(&format!("Sending message to [c:{:03}, wakeup:{}]: {:?}", self.id.0, try_wake_up, message));

        self.inbox.push(message);

        if try_wake_up {

            wake_up_if_sleeping(sched_ctx, self.id, self);

        }

    }

    fn increment_users(&self) {

        let old_count = self.num_handles.fetch_add(1, Ordering::AcqRel);

        debug_assert!(old_count > 0); // because we should never be able to retrieve a handle when the component is (being) destroyed

    }

    /// Returns the `CompKey` to the component if it should be destroyed

    pub(crate) fn decrement_users(&mut self) -> Option<CompKey> {

        debug_assert!(!self.decremented, "illegal to 'decrement_users' twice");

        let old_count = self.num_handles.fetch_sub(1, Ordering::AcqRel);

        let new_count = old_count - 1;

        dbg_code!(self.decremented = true);

        if new_count == 0 {

            return Some(unsafe{ self.id.upgrade() });

        }

        return None;

    }

}

impl Clone for CompHandle {

    fn clone(&self) -> Self {

        debug_assert!(!self.decremented, "illegal to clone after 'decrement_users'");

        self.increment_users();

        return CompHandle{

            target: self.target,

            id: self.id,

            #[cfg(debug_assertions)] decremented: false,

        };

    }

}

impl std::ops::Deref for CompHandle {

    type Target = CompPublic;

    fn deref(&self) -> &Self::Target {

        debug_assert!(!self.decremented); // cannot access if control is relinquished

        return unsafe{ &*self.target };

    }

}

impl Drop for CompHandle {

    fn drop(&mut self) {

        debug_assert!(self.decremented, "need call to 'decrement_users' before dropping");

    }

}

// -----------------------------------------------------------------------------

// Runtime

// -----------------------------------------------------------------------------

pub struct Runtime {

    pub(crate) inner: Arc<RuntimeInner>,

    threads: Vec<std::thread::JoinHandle<()>>,

}

impl Runtime {

    pub fn new(num_threads: u32, protocol_description: ProtocolDescription) -> Runtime {

        assert!(num_threads > 0, "need a thread to perform work");

        let runtime_inner = Arc::new(RuntimeInner {

            protocol: protocol_description,

            components: ComponentStore::new(128),

            work_queue: Mutex::new(VecDeque::with_capacity(128)),

            work_condvar: Condvar::new(),

            active_elements: AtomicU32::new(1),

        });

        let mut runtime = Runtime {

            inner: runtime_inner,

            threads: Vec::with_capacity(num_threads as usize),

        };

        for thread_index in 0..num_threads {

            let mut scheduler = Scheduler::new(runtime.inner.clone(), thread_index);

            let thread_handle = std::thread::spawn(move || {

                scheduler.run();

            });

            runtime.threads.push(thread_handle);

        }

        return runtime;

    }

    pub fn create_component(&self, module_name: &[u8], routine_name: &[u8]) -> Result<(), ComponentCreationError> {

        use crate::protocol::eval::ValueGroup;

        let prompt = self.inner.protocol.new_component(

            module_name, routine_name,

            ValueGroup::new_stack(Vec::new())

        )?;

        let reserved = self.inner.start_create_pdl_component();

        let ctx = CompCtx::new(&reserved);

        let (key, _) = self.inner.finish_create_pdl_component(reserved, CompPDL::new(prompt, 0), ctx, false);

        self.inner.enqueue_work(key);

        return Ok(())

    }

}

impl Drop for Runtime {

    fn drop(&mut self) {

        self.inner.decrement_active_components();

        for handle in self.threads.drain(..) {

            handle.join().expect("join scheduler thread");

        }

    }

}

/// Memory that is maintained by "the runtime". In practice it is maintained by

/// multiple schedulers, and this serves as the common interface to that memory.

pub(crate) struct RuntimeInner {

    pub protocol: ProtocolDescription,

    components: ComponentStore<RuntimeComp>,

    work_queue: Mutex<VecDeque<CompKey>>,

    work_condvar: Condvar,

    active_elements: AtomicU32, // active components and APIs (i.e. component creators)

}

impl RuntimeInner {

    // Scheduling and retrieving work

    pub(crate) fn take_work(&self) -> Option<CompKey> {

        let mut lock = self.work_queue.lock().unwrap();

        while lock.is_empty() && self.active_elements.load(Ordering::Acquire) != 0 {

            lock = self.work_condvar.wait(lock).unwrap();

        }

        // We have work, or the schedulers should exit.

        return lock.pop_front();

    }

    pub(crate) fn enqueue_work(&self, key: CompKey) {

        let mut lock = self.work_queue.lock().unwrap();

        lock.push_back(key);

        self.work_condvar.notify_one();

    }

    // Creating/destroying components

    pub(crate) fn start_create_pdl_component(&self) -> CompReserved {

        self.increment_active_components();

        let reservation = self.components.reserve();

        return CompReserved{ reservation };

    }

    pub(crate) fn finish_create_pdl_component(

        &self, reserved: CompReserved,

        component: CompPDL, mut context: CompCtx, initially_sleeping: bool,

    ) -> (CompKey, &mut RuntimeComp) {

        let inbox_queue = QueueDynMpsc::new(16);

        let inbox_producer = inbox_queue.producer();

        let _id = reserved.id();

        context.id = reserved.id();

        let component = RuntimeComp {

            public: CompPublic{

                sleeping: AtomicBool::new(initially_sleeping),

                num_handles: AtomicU32::new(1), // the component itself acts like a handle

                inbox: inbox_producer,

            },

            code: component,

            ctx: context,

            inbox: inbox_queue,

            exiting: false,

        };

        let index = self.components.submit(reserved.reservation, component);

        debug_assert_eq!(index, _id.0);

        let component = self.components.get_mut(index);

        return (CompKey(index), component);

    }

    pub(crate) fn get_component(&self, key: CompKey) -> &mut RuntimeComp {

        let component = self.components.get_mut(key.0);

        return component;

    }

    pub(crate) fn get_component_public(&self, id: CompId) -> CompHandle {

        let component = self.components.get(id.0);

        return CompHandle::new(id, &component.public);

    }

    pub(crate) fn destroy_component(&self, key: CompKey) {

        dbg_code!({

            let component = self.get_component(key);

            debug_assert!(component.exiting);

            debug_assert_eq!(component.public.num_handles.load(Ordering::Acquire), 0);

        });

        self.decrement_active_components();

        self.components.destroy(key.0);

    }

    // Tracking number of active interfaces and the active components

    #[inline]

    fn increment_active_components(&self) {

        let _old_val = self.active_elements.fetch_add(1, Ordering::AcqRel);

        debug_assert!(_old_val > 0); // can only create a component from a API/component, so can never be 0.

    }

    fn decrement_active_components(&self) {

        let old_val = self.active_elements.fetch_sub(1, Ordering::AcqRel);

        debug_assert!(old_val > 0); // something wrong with incr/decr logic

        let new_val = old_val - 1;

        if new_val == 0 {

            // Just to be sure, in case the last thing that gets destroyed is an

            // API instead of a thread.

            let _lock = self.work_queue.lock();

            self.work_condvar.notify_all();

        }

    }

}