use std::fs::File;

use std::io::Read;

use clap::{App, Arg};

use reowolf_rs as rw;

fn main() {

    let app = App::new("rwc")

        .author("Henger, M.")

        .version(env!("CARGO_PKG_VERSION"))

        .about("Reowolf compiler")

        .arg(

            Arg::new("input")

                .long("input")

                .short('i')

                .help("input files")

                .required(true)

                .takes_value(true)

                .multiple_occurrences(true)

        )

        .arg(

            Arg::new("threads")

                .long("threads")

                .short('t')

                .help("number of runtime threads")

                .default_value("1")

                .takes_value(true)

        );

    // Retrieve arguments and convert

    let app = app.get_matches();

    let input_files = app.values_of("input");

    if input_files.is_none() {

        println!("ERROR: Expected at least one input file");

        return;

    }

    let num_threads = app.value_of("threads").unwrap();

    let num_threads = match num_threads.parse::<i32>() {

        Ok(num_threads) => {

            if num_threads < 0 || num_threads > 255 {

                println!("ERROR: Number of threads must be a number between 0 and 256");

                return;

            }

            num_threads as u32

        },

        Err(err) => {

            println!("ERROR: Failed to parse number of threads\nbecause: {}", err);

            return;

        }

    };

    // Add input files to file buffer

    let input_files = input_files.unwrap();

    assert!(input_files.len() > 0); // because arg is required

    let mut builder = rw::ProtocolDescriptionBuilder::new();

    let mut file_buffer = Vec::with_capacity(4096);

    for input_file in input_files {

        print!("Adding file: {} ... ", input_file);

        let mut file = match File::open(input_file) {

            Ok(file) => file,

            Err(err) => {

                println!("FAILED (to open file)\nbecause:\n{}", err);

                return;

            }

        };

        file_buffer.clear();

        if let Err(err) = file.read_to_end(&mut file_buffer) {

            println!("FAILED (to read file)\nbecause:\n{}", err);

            return;

        }

        if let Err(err) = builder.add(input_file.to_string(), file_buffer.clone()) {

            println!("FAILED (to tokenize file)\nbecause:\n{}", err);

        }

        println!("Success");

    }

    // Compile the program

    print!("Compiling program ... ");

    let protocol_description = match builder.compile() {

        Ok(pd) => pd,

        Err(err) => {

            println!("FAILED\nbecause:\n{}", err);

            return;

        }

    };

    println!("Success");

    // Make sure there is a nameless module with a main component

    print!("Creating main component ... ");

    if let Err(err) = runtime.create_component(b"", b"main") {

        use rw::ComponentCreationError as CCE;

        let reason = match err {

            CCE::ModuleDoesntExist => "Input files did not contain a nameless module (that should contain the 'main' component)",

            CCE::DefinitionDoesntExist => "Input files did not contain a component called 'main'",

            CCE::DefinitionNotComponent => "Input file contained a 'main' function, but not a 'main' component",

            _ => "Unexpected error"

        };

        println!("FAILED\nbecause:\n{} (raw error: {:?})", reason, err);

        return;

    }

    println!("Success");

    println!("Now running until all components have exited");

    println!("--------------------------------------------\n\n");

}

    fn handle_unblock_port_instruction(&mut self, sched_ctx: &SchedulerCtx, comp_ctx: &mut CompCtx, port_handle: LocalPortHandle) {

        let port_info = comp_ctx.get_port_mut(port_handle);

        let port_id = port_info.self_id;

        debug_assert!(port_info.state.is_blocked());

        port_info.state = PortState::Open;

        if self.mode == Mode::BlockedPut && port_id == self.mode_port {

            // We were blocked on the port that just became unblocked, so

            // send the message.

            debug_assert_eq!(port_info.kind, PortKind::Putter);

            let mut replacement = ValueGroup::default();

            std::mem::swap(&mut replacement, &mut self.mode_value);

            self.send_data_message_and_wake_up(sched_ctx, comp_ctx, port_handle, replacement);

            self.mode = Mode::Sync;

            self.mode_port = PortId::new_invalid();

        }

    }

    fn create_component_and_transfer_ports(

        &mut self,

        sched_ctx: &SchedulerCtx, creator_ctx: &mut CompCtx,

        definition_id: DefinitionId, monomorph_index: i32, mut arguments: ValueGroup

    ) {

        struct PortPair{

            creator_handle: LocalPortHandle,

            creator_id: PortId,

            created_handle: LocalPortHandle,

            created_id: PortId,

        }

        let mut port_id_pairs = Vec::new();

        let reservation = sched_ctx.runtime.start_create_pdl_component();

        let mut created_ctx = CompCtx::new(&reservation);

        // Take all the ports ID that are in the `args` (and currently belong to

        // the creator component) and translate them into new IDs that are

        // associated with the component we're about to create

        let mut arg_iter = ValueGroupIter::new(&mut arguments);

        while let Some(port_reference) = arg_iter.next() {

            // Create port entry for new component

            let creator_port_id = port_reference.id;

            let creator_port_handle = creator_ctx.get_port_handle(creator_port_id);

            let creator_port = creator_ctx.get_port(creator_port_handle);

            let created_port_handle = created_ctx.add_port(

                creator_port.peer_comp_id, creator_port.peer_port_id,

                creator_port.kind, creator_port.state

            );

            let created_port = created_ctx.get_port(created_port_handle);

            let created_port_id = created_port.self_id;

            port_id_pairs.push(PortPair{

                creator_handle: creator_port_handle,

                creator_id: creator_port_id,

                created_handle: created_port_handle,

                created_id: created_port_id,

            });

            // Modify value in arguments (bit dirty, but double vec in ValueGroup causes lifetime issues)

            let arg_value = if let Some(heap_pos) = port_reference.heap_pos {

                &mut arg_iter.group.regions[heap_pos][port_reference.index]

            } else {

                &mut arg_iter.group.values[port_reference.index]

            };

            match arg_value {

                Value::Input(id) => *id = port_id_to_eval(created_port_id),

                Value::Output(id) => *id = port_id_to_eval(created_port_id),

                _ => unreachable!(),

            }

        }

        // For each transferred port pair set their peer components to the

        // correct values. This will only change the values for the ports of

        // the new component.

        let mut created_component_has_remote_peers = false;

        for pair in port_id_pairs.iter() {

            let creator_port_info = creator_ctx.get_port(pair.creator_handle);

            let created_port_info = created_ctx.get_port_mut(pair.created_handle);

            if created_port_info.peer_comp_id == creator_ctx.id {

                // Port peer is owned by the creator as well

                let created_peer_port_index = port_id_pairs

                    .iter()

                    .position(|v| v.creator_id == creator_port_info.peer_port_id);

                match created_peer_port_index {

                    Some(created_peer_port_index) => {

                        // Peer port moved to the new component as well. So

                        // adjust IDs appropriately.

                        let peer_pair = &port_id_pairs[created_peer_port_index];

                        created_port_info.peer_port_id = peer_pair.created_id;

                        created_port_info.peer_comp_id = reservation.id();

                        todo!("either add 'self peer', or remove that idea from Ctx altogether")

                    },

                    None => {

                        // Peer port remains with creator component.

                        created_port_info.peer_comp_id = creator_ctx.id;

                        created_ctx.add_peer(pair.created_handle, sched_ctx, creator_ctx.id, None);

                    }

                }

            } else {

                // Peer is a different component. We'll deal with sending the

                // appropriate messages later

                let peer_handle = creator_ctx.get_peer_handle(created_port_info.peer_comp_id);

                let peer_info = creator_ctx.get_peer(peer_handle);

                created_ctx.add_peer(pair.created_handle, sched_ctx, peer_info.id, Some(&peer_info.handle));

                created_component_has_remote_peers = true;

            }

        }

        // We'll now actually turn our reservation for a new component into an

        // actual component. Note that we initialize it as "not sleeping" as

        // its initial scheduling might be performed based on `Ack`s in response

        // to message exchanges between remote peers.

        let prompt = Prompt::new(

            &sched_ctx.runtime.protocol.types, &sched_ctx.runtime.protocol.heap,

            definition_id, monomorph_index, arguments,

        );

        let component = CompPDL::new(prompt, port_id_pairs.len());

        let (created_key, component) = sched_ctx.runtime.finish_create_pdl_component(

            reservation, component, created_ctx, false,

        );

        let created_ctx = &component.ctx;

        // Now modify the creator's ports: remove every transferred port and

        // potentially remove the peer component. Here is also where we will

        // transfer messages in the main inbox.

        for pair in port_id_pairs.iter() {

            // Remove peer if appropriate

            let creator_port_info = creator_ctx.get_port(pair.creator_handle);

            let creator_port_index = creator_ctx.get_port_index(pair.creator_handle);

            let creator_peer_comp_id = creator_port_info.peer_comp_id;

            creator_ctx.remove_peer(sched_ctx, pair.creator_handle, creator_peer_comp_id, false);

            creator_ctx.remove_port(pair.creator_handle);

            // Transfer any messages

            let created_port_index = created_ctx.get_port_index(pair.created_handle);

            let created_port_info = created_ctx.get_port(pair.created_handle);

            debug_assert!(component.code.inbox_main[created_port_index].is_none());

            if let Some(mut message) = self.inbox_main.remove(creator_port_index) {

                message.data_header.target_port = pair.created_id;

                component.code.inbox_main[created_port_index] = Some(message);

            }

            let mut message_index = 0;

            while message_index < self.inbox_backup.len() {

                let message = &self.inbox_backup[message_index];

                if message.data_header.target_port == pair.creator_id {

                    // transfer message

                    let mut message = self.inbox_backup.remove(message_index);

                    message.data_header.target_port = pair.created_id;

                    component.code.inbox_backup.push(message);

                } else {

                    message_index += 1;

                }

            }

            // Handle potential channel between creator and created component

            if created_port_info.peer_comp_id == creator_ctx.id {

                let peer_port_handle = creator_ctx.get_port_handle(created_port_info.peer_port_id);

                let peer_port_info = creator_ctx.get_port_mut(peer_port_handle);

                peer_port_info.peer_comp_id = created_ctx.id;

                peer_port_info.peer_port_id = created_port_info.self_id;

                creator_ctx.add_peer(peer_port_handle, sched_ctx, created_ctx.id, None);

            }

        }

        // By now all ports have been transferred. We'll now do any of the setup

        // for rerouting/messaging

        if created_component_has_remote_peers {

            let schedule_entry_id = self.control.add_schedule_entry(created_ctx.id);

            for pair in port_id_pairs.iter() {

                let port_info = created_ctx.get_port(pair.created_handle);

                if port_info.peer_comp_id != creator_ctx.id && port_info.peer_comp_id != created_ctx.id {

                    let message = self.control.add_reroute_entry(

                        creator_ctx.id, port_info.peer_port_id, port_info.peer_comp_id,

                        pair.creator_id, pair.created_id, created_ctx.id,

                        schedule_entry_id

                    );

                    let peer_handle = created_ctx.get_peer_handle(port_info.peer_comp_id);

                    let peer_info = created_ctx.get_peer(peer_handle);

                    peer_info.handle.send_message(sched_ctx, message, true);

                }

            }

        } else {

            // Peer can be scheduled immediately

            sched_ctx.runtime.enqueue_work(created_key);

        }

    }

}

#[inline]

/// 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 {

        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 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);

use std::sync::Arc;

use std::sync::atomic::Ordering;

use super::component::*;

use super::runtime::*;

/// Data associated with a scheduler thread

pub(crate) struct Scheduler {

    runtime: Arc<RuntimeInner>,

    scheduler_id: u32,

}

pub(crate) struct SchedulerCtx<'a> {

    pub runtime: &'a RuntimeInner,

    pub id: u32,

    pub comp: u32,

}

impl<'a> SchedulerCtx<'a> {

        return Self {

            runtime,

            id,

            comp: 0,

        }

    }

    pub(crate) fn log(&self, text: &str) {

            println!("[s:{:02}, c:{:03}] {}", self.id, self.comp, text);

        }

    // TODO: Obviously remove, but useful for testing

    pub(crate) fn log_special(&self, text: &str) {

            println!("[s:{:02}, c:{:03}] *** *** {}", self.id, self.comp, text);

        }

    }

impl Scheduler {

    // public interface to thread

    }

    pub fn run(&mut self) {

        'run_loop: loop {

            // Wait until we have something to do (or need to quit)

            let comp_key = self.runtime.take_work();

            if comp_key.is_none() {

                break 'run_loop;

            }

            let comp_key = comp_key.unwrap();

            let component = self.runtime.get_component(comp_key);

            scheduler_ctx.comp = comp_key.0;

            // Run the component until it no longer indicates that it needs to

            // be re-executed immediately.

            let mut new_scheduling = CompScheduling::Immediate;

            while let CompScheduling::Immediate = new_scheduling {

                while let Some(message) = component.inbox.pop() {

                    component.code.handle_message(&mut scheduler_ctx, &mut component.ctx, message);

                }

                new_scheduling = component.code.run(&mut scheduler_ctx, &mut component.ctx).expect("TODO: Handle error");

            }

            // Handle the new scheduling

            match new_scheduling {

                CompScheduling::Immediate => unreachable!(),

                CompScheduling::Requeue => { self.runtime.enqueue_work(comp_key); },

                CompScheduling::Sleep => { self.mark_component_as_sleeping(comp_key, component); },

                CompScheduling::Exit => { self.mark_component_as_exiting(&scheduler_ctx, component); }

            }

        }

    }

    // local utilities

    /// Marks component as sleeping, if after marking itself as sleeping the

    /// inbox contains messages then the component will be immediately

    /// rescheduled. After calling this function the component should not be

    /// executed anymore.

    fn mark_component_as_sleeping(&self, key: CompKey, component: &mut RuntimeComp) {

        debug_assert_eq!(key.downgrade(), component.ctx.id); // make sure component matches key

        debug_assert_eq!(component.public.sleeping.load(Ordering::Acquire), false); // we're executing it, so it cannot be sleeping

        component.public.sleeping.store(true, Ordering::Release);

        if component.inbox.can_pop() {

            let should_reschedule = component.public.sleeping

                .compare_exchange(true, false, Ordering::AcqRel, Ordering::Relaxed)

                .is_ok();

            if should_reschedule {

                self.runtime.enqueue_work(key);

            }

        }

    }

    /// Marks the component as exiting by removing the reference it holds to

    /// itself. Afterward the component will enter "normal" sleeping mode (if it

    /// has not yet been destroyed)

    fn mark_component_as_exiting(&self, sched_ctx: &SchedulerCtx, component: &mut RuntimeComp) {

        // If we didn't yet decrement our reference count, do so now

        let comp_key = unsafe{ component.ctx.id.upgrade() };

        if !component.exiting {

            component.exiting = true;

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

            let new_count = old_count - 1;

            if new_count == 0 {

                sched_ctx.runtime.destroy_component(comp_key);

                return;

            }

        }

        // Enter "regular" sleeping mode

        self.mark_component_as_sleeping(comp_key, component);

    }

}

use crate::protocol::*;

use crate::protocol::eval::*;

use crate::runtime2::runtime::*;

use crate::runtime2::component::{CompCtx, CompPDL};

fn create_component(rt: &Runtime, module_name: &str, routine_name: &str, args: ValueGroup) {

    let prompt = rt.inner.protocol.new_component(

        module_name.as_bytes(), routine_name.as_bytes(), args

    ).expect("create prompt");

    let reserved = rt.inner.start_create_pdl_component();

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

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

    rt.inner.enqueue_work(key);

}

fn no_args() -> ValueGroup { ValueGroup::new_stack(Vec::new()) }

#[test]

fn test_component_creation() {

    let pd = ProtocolDescription::parse(b"

    primitive nothing_at_all() {

        s32 a = 5;

        auto b = 5 + a;

    }

    ").expect("compilation");

    for i in 0..20 {

        create_component(&rt, "", "nothing_at_all", no_args());

    }

}

#[test]

fn test_component_communication() {

    let pd = ProtocolDescription::parse(b"

    primitive sender(out<u32> o, u32 outside_loops, u32 inside_loops) {

        u32 outside_index = 0;

        while (outside_index < outside_loops) {

            u32 inside_index = 0;

            sync while (inside_index < inside_loops) {

                put(o, inside_index);

                inside_index += 1;

            }

            outside_index += 1;

        }

    }

    primitive receiver(in<u32> i, u32 outside_loops, u32 inside_loops) {

        u32 outside_index = 0;

        while (outside_index < outside_loops) {

            u32 inside_index = 0;

            sync while (inside_index < inside_loops) {

                auto val = get(i);

                while (val != inside_index) {} // infinite loop if incorrect value is received

                inside_index += 1;

            }

            outside_index += 1;

        }

    }

    composite constructor() {

        channel o_orom -> i_orom;

        channel o_mrom -> i_mrom;

        channel o_ormm -> i_ormm;

        channel o_mrmm -> i_mrmm;

        // one round, one message per round

        new sender(o_orom, 1, 1);

        new receiver(i_orom, 1, 1);

        // multiple rounds, one message per round

        new sender(o_mrom, 5, 1);

        new receiver(i_mrom, 5, 1);

        // one round, multiple messages per round

        new sender(o_ormm, 1, 5);

        new receiver(i_ormm, 1, 5);

        // multiple rounds, multiple messages per round

        new sender(o_mrmm, 5, 5);

        new receiver(i_mrmm, 5, 5);

    }").expect("compilation");

    create_component(&rt, "", "constructor", no_args());

}

#module producer

primitive producer(out<u32> output) {

    sync {

        print("P: Going to send a value!");

        put(output, 1337);

        print("P: I just sent a value!");

    }

    print("P: I am exiting");

}

primitive consumer(in<u32> input) {

    sync {

        print("C: going to receive a value");

        auto v = get(input);

        print("C: I have received a value");

    }

    print("C: I am now exiting");

}

primitive producer(out<u32> output) {

    sync {

        print("P: Going to send a value!");

        put(output, 1337);

        print("P: I just sent a value!");

    }

    print("P: I am exiting");

}

composite main() {

    channel output -> input;

    new consumer(input);

    new producer(output);

}