mod arena;
pub(crate) mod eval;
pub(crate) mod input_source;
mod parser;
#[cfg(test)] mod tests;

pub(crate) mod ast;
pub(crate) mod ast_printer;

use std::sync::Mutex;

use crate::collections::{StringPool, StringRef};
use crate::common::*;
use crate::protocol::ast::*;
use crate::protocol::eval::*;
use crate::protocol::input_source::*;
use crate::protocol::parser::*;
use crate::protocol::type_table::*;

/// A protocol description module
pub struct Module {
    pub(crate) source: InputSource,
    pub(crate) root_id: RootId,
    pub(crate) name: Option<StringRef<'static>>,
}
/// Description of a protocol object, used to configure new connectors.
#[repr(C)]
pub struct ProtocolDescription {
    pub(crate) modules: Vec<Module>,
    pub(crate) heap: Heap,
    pub(crate) types: TypeTable,
    pub(crate) pool: Mutex<StringPool>,
}
#[derive(Debug, Clone)]
pub(crate) struct ComponentState {
    pub(crate) prompt: Prompt,
}

#[allow(dead_code)]
pub(crate) enum EvalContext<'a> {
    Nonsync(&'a mut NonsyncProtoContext<'a>),
    Sync(&'a mut SyncProtoContext<'a>),
    None,
}
//////////////////////////////////////////////

#[derive(Debug)]
pub enum ComponentCreationError {
    ModuleDoesntExist,
    DefinitionDoesntExist,
    DefinitionNotComponent,
    InvalidNumArguments,
    InvalidArgumentType(usize),
}

impl std::fmt::Debug for ProtocolDescription {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "(An opaque protocol description)")
    }
}
impl ProtocolDescription {
    // TODO: Allow for multi-file compilation
    pub fn parse(buffer: &[u8]) -> Result<Self, String> {
        // TODO: @fixme, keep code compilable, but needs support for multiple
        //  input files.
        let source = InputSource::new(String::new(), Vec::from(buffer));
        let mut parser = Parser::new();
        parser.feed(source).expect("failed to feed source");
        
        if let Err(err) = parser.parse() {
            println!("ERROR:\n{}", err);
            return Err(format!("{}", err))
        }

        debug_assert_eq!(parser.modules.len(), 1, "only supporting one module here for now");
        let modules: Vec<Module> = parser.modules.into_iter()
            .map(|module| Module{
                source: module.source,
                root_id: module.root_id,
                name: module.name.map(|(_, name)| name)
            })
            .collect();

        return Ok(ProtocolDescription {
            modules,
            heap: parser.heap,
            types: parser.type_table,
            pool: Mutex::new(parser.string_pool),
        });
    }

    #[deprecated]
    pub(crate) fn component_polarities(
        &self,
        module_name: &[u8],
        identifier: &[u8],
    ) -> Result<Vec<Polarity>, AddComponentError> {
        use AddComponentError::*;

        let module_root = self.lookup_module_root(module_name);
        if module_root.is_none() {
            return Err(AddComponentError::NoSuchModule);
        }
        let module_root = module_root.unwrap();

        let root = &self.heap[module_root];
        let def = root.get_definition_ident(&self.heap, identifier);
        if def.is_none() {
            return Err(NoSuchComponent);
        }

        let def = &self.heap[def.unwrap()];
        if !def.is_component() {
            return Err(NoSuchComponent);
        }

        for &param in def.parameters().iter() {
            let param = &self.heap[param];
            let first_element = &param.parser_type.elements[0];

            match first_element.variant {
                ParserTypeVariant::Input | ParserTypeVariant::Output => continue,
                _ => {
                    return Err(NonPortTypeParameters);
                }
            }
        }

        let mut result = Vec::new();
        for &param in def.parameters().iter() {
            let param = &self.heap[param];
            let first_element = &param.parser_type.elements[0];

            if first_element.variant == ParserTypeVariant::Input {
                result.push(Polarity::Getter)
            } else if first_element.variant == ParserTypeVariant::Output {
                result.push(Polarity::Putter)
            } else {
                unreachable!()
            }
        }
        Ok(result)
    }

    // expects port polarities to be correct
    #[deprecated]
    pub(crate) fn new_component(&self, module_name: &[u8], identifier: &[u8], ports: &[PortId]) -> ComponentState {
        let mut args = Vec::new();
        for (&x, y) in ports.iter().zip(self.component_polarities(module_name, identifier).unwrap()) {
            match y {
                Polarity::Getter => args.push(Value::Input(x)),
                Polarity::Putter => args.push(Value::Output(x)),
            }
        }

        let module_root = self.lookup_module_root(module_name).unwrap();
        let root = &self.heap[module_root];
        let def = root.get_definition_ident(&self.heap, identifier).unwrap();
        // TODO: Check for polymorph
        ComponentState { prompt: Prompt::new(&self.types, &self.heap, def, 0, ValueGroup::new_stack(args)) }
    }

    // TODO: Ofcourse, rename this at some point, perhaps even remove it in its
    //  entirety. Find some way to interface with the parameter's types.
    pub(crate) fn new_component_v2(
        &self, module_name: &[u8], identifier: &[u8], arguments: ValueGroup
    ) -> Result<ComponentState, ComponentCreationError> {
        // Find the module in which the definition can be found
        let module_root = self.lookup_module_root(module_name);
        if module_root.is_none() {
            return Err(ComponentCreationError::ModuleDoesntExist);
        }
        let module_root = module_root.unwrap();

        let root = &self.heap[module_root];
        let definition_id = root.get_definition_ident(&heap, identifier);
        if definition_id.is_none() {
            return Err(ComponentCreationError::DefinitionDoesntExist);
        }
        let definition_id = definition_id.unwrap();

        let definition = &self.heap[definition_id];
        if !definition.is_component() {
            return Err(ComponentCreationError::DefinitionNotComponent);
        }

        // Make sure that the types of the provided value group matches that of
        // the expected types.
        let definition = definition.as_component();
        if !definition.poly_vars.is_empty() {
            return Err(ComponentCreationError::DefinitionNotComponent);
        }

        // - check number of arguments
        let expr_data = self.types.get_procedure_expression_data(&definition_id, 0);
        if expr_data.arg_types.len() != arguments.values.len() {
            return Err(ComponentCreationError::InvalidNumArguments);
        }

        // - for each argument try to make sure the types match
        for arg_idx in 0..arguments.values.len() {
            let expected_type = &expr_data.arg_types[arg_idx];
            let provided_value = &arguments.values[arg_idx];
            if !self.verify_same_type(expected_type, 0, &arguments, provided_value) {
                return Err(ComponentCreationError::InvalidArgumentType(arg_idx));
            }
        }

        // By now we're sure that all of the arguments are correct. So create
        // the connector.
        return Ok(ComponentState{
            prompt: Prompt::new(&self.types, &self.heap, def, 0, arguments),
        });
    }

    fn lookup_module_root(&self, module_name: &[u8]) -> Option<RootId> {
        for module in self.modules.iter() {
            match &module.name {
                Some(name) => if name.as_bytes() == module_name {
                    return Some(module.root_id);
                },
                None => if module_name.is_empty() {
                    return Some(module.root_id);
                }
            }
        }

        return None;
    }

    fn verify_same_type(&self, expected: &ConcreteType, expected_idx: usize, arguments: &ValueGroup, argument: &Value) -> bool {
        use ConcreteTypePart as CTP;

        macro_rules! match_variant {
            ($value:expr, $variant:expr) => {
                if let $variant(_) = $value { true } else { false }
            };
        }

        match &expected.parts[expected_idx] {
            CTP::Void | CTP::Message | CTP::Slice | CTP::Function(_, _) | CTP::Component(_, _) => unreachable!(),
            CTP::Bool => match_variant!(argument, Value::Bool),
            CTP::UInt8 => match_variant!(argument, Value::UInt8),
            CTP::UInt16 => match_variant!(argument, Value::UInt16),
            CTP::UInt32 => match_variant!(argument, Value::UInt32),
            CTP::UInt64 => match_variant!(argument, Value::UInt64),
            CTP::SInt8 => match_variant!(argument, Value::SInt8),
            CTP::SInt16 => match_variant!(argument, Value::SInt16),
            CTP::SInt32 => match_variant!(argument, Value::SInt32),
            CTP::SInt64 => match_variant!(argument, Value::SInt64),
            CTP::Character => match_variant!(argument, Value::Char),
            CTP::String => {
                // Match outer string type and embedded character types
                if let Value::String(heap_pos) = argument {
                    for element in &arguments.regions[*heap_pos as usize] {
                        if let Value::Char(_) = element {} else {
                            return false;
                        }
                    }
                } else {
                    return false;
                }

                return true;
            },
            CTP::Array => {
                if let Value::Array(heap_pos) = argument {
                    let heap_pos = *heap_pos;
                    for element in &arguments.regions[heap_pos as usize] {
                        if !self.verify_same_type(expected, expected_idx + 1, arguments, element) {
                            return false;
                        }
                    }
                    return true;
                } else {
                    return false;
                }
            },
            CTP::Input => match_variant!(argument, Value::Input),
            CTP::Output => match_variant!(argument, Value::Output),
            CTP::Instance(_definition_id, _num_embedded) => {
                todo!("implement full type checking on user-supplied arguments");
                return false;
            },
        }
    }
}

// TODO: @temp Should just become a concrete thing that is passed in
pub trait RunContext {
    fn did_put(&mut self, port: PortId) -> bool;
    fn get(&mut self, port: PortId) -> Option<ValueGroup>; // None if still waiting on message
    fn fires(&mut self, port: PortId) -> Option<Value>; // None if not yet branched
    fn get_channel(&mut self) -> Option<(Value, Value)>; // None if not yet prepared
}

#[derive(Debug)]
pub enum RunResult {
    // Can only occur outside sync blocks
    ComponentTerminated, // component has exited its procedure
    ComponentAtSyncStart,
    NewComponent(DefinitionId, i32, ValueGroup), // should also be possible inside sync
    NewChannel, // should also be possible inside sync
    // Can only occur inside sync blocks
    BranchInconsistent, // branch has inconsistent behaviour
    BranchMissingPortState(PortId), // branch doesn't know about port firing
    BranchMissingPortValue(PortId), // branch hasn't received message on input port yet
    BranchAtSyncEnd,
    BranchPut(PortId, ValueGroup),
}

impl ComponentState {
    pub(crate) fn run(&mut self, ctx: &mut impl RunContext, pd: &ProtocolDescription) -> RunResult {
        use EvalContinuation as EC;
        use RunResult as RR;

        loop {
            let step_result = self.prompt.step(&pd.types, &pd.heap, &pd.modules, ctx);
            match step_result {
                Err(reason) => {
                    // TODO: @temp
                    println!("Evaluation error:\n{}", reason);
                    todo!("proper error handling/bubbling up");
                },
                Ok(continuation) => match continuation {
                    // TODO: Probably want to remove this translation
                    EC::Stepping => continue,
                    EC::Inconsistent => return RR::BranchInconsistent,
                    EC::Terminal => return RR::ComponentTerminated,
                    EC::SyncBlockStart => return RR::ComponentAtSyncStart,
                    EC::SyncBlockEnd => return RR::BranchAtSyncEnd,
                    EC::NewComponent(definition_id, monomorph_idx, args) =>
                        return RR::NewComponent(definition_id, monomorph_idx, args),
                    EC::NewChannel =>
                        return RR::NewChannel,
                    EC::BlockFires(port_id) => return RR::BranchMissingPortState(port_id),
                    EC::BlockGet(port_id) => return RR::BranchMissingPortValue(port_id),
                    EC::Put(port_id, value) => {
                        let value_group = ValueGroup::from_store(&self.prompt.store, &[value]);
                        return RR::BranchPut(port_id, value_group);
                    },
                }
            }
        }
    }
}

// TODO: @remove the old stuff
impl ComponentState {
    pub(crate) fn nonsync_run<'a: 'b, 'b>(
        &'a mut self,
        context: &'b mut NonsyncProtoContext<'b>,
        pd: &'a ProtocolDescription,
    ) -> NonsyncBlocker {
        let mut context = EvalContext::Nonsync(context);
        loop {
            let result = self.prompt.step(&pd.types, &pd.heap, &pd.modules, &mut context);
            match result {
                Err(err) => {
                    println!("Evaluation error:\n{}", err);
                    panic!("proper error handling when component fails");
                },
                Ok(cont) => match cont {
                    EvalContinuation::Stepping => continue,
                    EvalContinuation::Inconsistent => return NonsyncBlocker::Inconsistent,
                    EvalContinuation::Terminal => return NonsyncBlocker::ComponentExit,
                    EvalContinuation::SyncBlockStart => return NonsyncBlocker::SyncBlockStart,
                    // Not possible to end sync block if never entered one
                    EvalContinuation::SyncBlockEnd => unreachable!(),
                    EvalContinuation::NewComponent(definition_id, monomorph_idx, args) => {
                        // Look up definition (TODO for now, assume it is a definition)
                        let mut moved_ports = HashSet::new();
                        for arg in args.values.iter() {
                            match arg {
                                Value::Output(port) => {
                                    moved_ports.insert(*port);
                                }
                                Value::Input(port) => {
                                    moved_ports.insert(*port);
                                }
                                _ => {}
                            }
                        }
                        for region in args.regions.iter() {
                            for arg in region {
                                match arg {
                                    Value::Output(port) => { moved_ports.insert(*port); },
                                    Value::Input(port) => { moved_ports.insert(*port); },
                                    _ => {},
                                }
                            }
                        }
                        let init_state = ComponentState { prompt: Prompt::new(&pd.types, &pd.heap, definition_id, monomorph_idx, args) };
                        context.new_component(moved_ports, init_state);
                        // Continue stepping
                        continue;
                    },
                    EvalContinuation::NewChannel => {
                        // Because of the way we emulate the old context for now, we can safely
                        // assume that this will never happen. The old context thingamajig always
                        // creates a channel, it never bubbles a "need to create a channel" message
                        // to the runtime
                        unreachable!();
                    },
                    // Outside synchronous blocks, no fires/get/put happens
                    EvalContinuation::BlockFires(_) => unreachable!(),
                    EvalContinuation::BlockGet(_) => unreachable!(),
                    EvalContinuation::Put(_, _) => unreachable!(),
                },
            }
        }
    }

    pub(crate) fn sync_run<'a: 'b, 'b>(
        &'a mut self,
        context: &'b mut SyncProtoContext<'b>,
        pd: &'a ProtocolDescription,
    ) -> SyncBlocker {
        let mut context = EvalContext::Sync(context);
        loop {
            let result = self.prompt.step(&pd.types, &pd.heap, &pd.modules, &mut context);
            match result {
                Err(err) => {
                    println!("Evaluation error:\n{}", err);
                    panic!("proper error handling when component fails");
                },
                Ok(cont) => match cont {
                    EvalContinuation::Stepping => continue,
                    EvalContinuation::Inconsistent => return SyncBlocker::Inconsistent,
                    // First need to exit synchronous block before definition may end
                    EvalContinuation::Terminal => unreachable!(),
                    // No nested synchronous blocks
                    EvalContinuation::SyncBlockStart => unreachable!(),
                    EvalContinuation::SyncBlockEnd => return SyncBlocker::SyncBlockEnd,
                    // Not possible to create component in sync block
                    EvalContinuation::NewComponent(_, _, _) => unreachable!(),
                    EvalContinuation::NewChannel => unreachable!(),
                    EvalContinuation::BlockFires(port) => {
                        return SyncBlocker::CouldntCheckFiring(port);
                    },
                    EvalContinuation::BlockGet(port) => {
                        return SyncBlocker::CouldntReadMsg(port);
                    },
                    EvalContinuation::Put(port, message) => {
                        let payload;
                        match message {
                            Value::Null => {
                                return SyncBlocker::Inconsistent;
                            },
                            Value::Message(heap_pos) => {
                                // Create a copy of the payload
                                let values = &self.prompt.store.heap_regions[heap_pos as usize].values;
                                let mut bytes = Vec::with_capacity(values.len());
                                for value in values {
                                    bytes.push(value.as_uint8());
                                }
                                payload = Payload(Arc::new(bytes));
                            }
                            _ => unreachable!(),
                        }
                        return SyncBlocker::PutMsg(port, payload);
                    }
                },
            }
        }
    }
}

impl RunContext for EvalContext<'_> {
    fn did_put(&mut self, port: PortId) -> bool {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(_) => unreachable!(),
            EvalContext::Sync(ctx) => {
                ctx.did_put_or_get(port)
            }
        }
    }

    fn get(&mut self, port: PortId) -> Option<ValueGroup> {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(_) => unreachable!(),
            EvalContext::Sync(ctx) => {
                let payload = ctx.read_msg(port);
                if payload.is_none() {
                    return None;
                }

                let payload = payload.unwrap();
                let mut transformed = Vec::with_capacity(payload.len());
                for byte in payload.0.iter() {
                    transformed.push(Value::UInt8(*byte));
                }

                let value_group = ValueGroup{
                    values: vec![Value::Message(0)],
                    regions: vec![transformed],
                };

                return Some(value_group);
            }
        }
    }

    fn fires(&mut self, port: PortId) -> Option<Value> {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(_) => unreachable!(),
            EvalContext::Sync(context) => {
                match context.is_firing(port) {
                    Some(did_fire) => Some(Value::Bool(did_fire)),
                    None => None,
                }
            }
        }
    }

    fn get_channel(&mut self) -> Option<(Value, Value)> {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(context) => {
                let [from, to] = context.new_port_pair();
                let from = Value::Output(from);
                let to = Value::Input(to);
                return Some((from, to));
            },
            EvalContext::Sync(_) => unreachable!(),
        }
    }
}

// TODO: @remove once old runtime has disappeared
impl EvalContext<'_> {
    // fn random(&mut self) -> LongValue {
    //     match self {
    //         // EvalContext::None => unreachable!(),
    //         EvalContext::Nonsync(_context) => todo!(),
    //         EvalContext::Sync(_) => unreachable!(),
    //     }
    // }
    fn new_component(&mut self, moved_ports: HashSet<PortId>, init_state: ComponentState) -> () {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(context) => {
                context.new_component(moved_ports, init_state)
            }
            EvalContext::Sync(_) => unreachable!(),
        }
    }
    fn new_channel(&mut self) -> [Value; 2] {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(context) => {
                let [from, to] = context.new_port_pair();
                let from = Value::Output(from);
                let to = Value::Input(to);
                return [from, to];
            }
            EvalContext::Sync(_) => unreachable!(),
        }
    }
    fn fires(&mut self, port: Value) -> Option<Value> {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(_) => unreachable!(),
            EvalContext::Sync(context) => match port {
                Value::Output(port) => context.is_firing(port).map(Value::Bool),
                Value::Input(port) => context.is_firing(port).map(Value::Bool),
                _ => unreachable!(),
            },
        }
    }
    fn get(&mut self, port: Value, store: &mut Store) -> Option<Value> {
        match self {
            EvalContext::None => unreachable!(),
            EvalContext::Nonsync(_) => unreachable!(),
            EvalContext::Sync(context) => match port {
                Value::Input(port) => {
                    let payload = context.read_msg(port);
                    if payload.is_none() { return None; }

                    let heap_pos = store.alloc_heap();
                    let heap_pos_usize = heap_pos as usize;
                    let payload = payload.unwrap();
                    store.heap_regions[heap_pos_usize].values.reserve(payload.0.len());
                    for value in payload.0.iter() {
                        store.heap_regions[heap_pos_usize].values.push(Value::UInt8(*value));
                    }

                    return Some(Value::Message(heap_pos));
                }
                _ => unreachable!(),
            },
        }
    }
    fn did_put(&mut self, port: Value) -> bool {
        match self {
            EvalContext::None => unreachable!("did_put in None context"),
            EvalContext::Nonsync(_) => unreachable!("did_put in nonsync context"),
            EvalContext::Sync(context) => match port {
                Value::Output(port) => {
                    context.did_put_or_get(port)
                },
                _ => unreachable!("did_put on non-output port value")
            }
        }
    }
}