//! Squad of cooks ready to cook your data.
//!
//! You can use cooks directly on `Vec` or `Iterator` via `Cookable` trait. Also
//! you can reuse once created cooks on `Iterator` and `Vec` via `Cooks::cook`
//! and `Cooks::cook_vec`.
//!
//! The case of losing data is when cook panicking while cooking it.
//!
//! The case of main thread to panic is when it successfully `recv` cook's result
//! but could not `send` him a new input, which, i think, is almost impossible.
//! (Don't look like that, Mr. Murphy!).
//!
//! #### `Vec`
//!
//! You can call `cook` on `Vec<Input>` and it will block until you get
//! `Vec<Option<Output>>` of cooked data in the same order.
//!
//! In case of panicking cook he will be revived but you will get `None` in
//! place of it's data.
//!
//! #### `Iterator`
//!
//! You can call `cook` on `Iterator<Input>` and you will get `Dishes`, which is
//! `Iterator<Output>`. Call of `next` on it will block until some cook finished
//! cooking. There is no guarantees about the order of returned data.
//!
//! Panicking cooks will be revived, but data will be lost.
#![feature(unsafe_destructor, std_misc, os, core)]
use std::sync::mpsc::{Select, SyncSender, Receiver, sync_channel, SendError};
use std::iter::{IteratorExt};
use std::thread::Thread;
use std::sync::Arc;

use self::State::*;
use self::Value::*;

#[derive(PartialEq, Eq, Clone)]
enum State {
    Idling,
    Working,
    Done,
}

enum Value<T> {
    Indexed(usize, T),
    Plain(T),
}

pub struct Cooks<Txty, Rxty, F> {
    txs: Vec<SyncSender<Value<Option<Txty>>>>,
    rxs: Vec<Receiver<Value<Option<Rxty>>>>,
    states: Vec<State>,
    f: Arc<F>,
}

impl<Txty: Send, Rxty: Send, F> Cooks<Txty, Rxty, F>
where F: Fn(Txty) -> Option<Rxty>, F: Send + Sync + 'static {

    /// Starts new `num_cpus` cooks.
    pub fn new(f: F) -> Cooks<Txty, Rxty, F> {
        Cooks::new_n(::std::os::num_cpus(), f)
    }

    /// Starts new `n` cooks.
    pub fn new_n(count: usize, f: F) -> Cooks<Txty, Rxty, F> {
        let mut cooks = Cooks {
            txs: Vec::with_capacity(count),
            rxs: Vec::with_capacity(count),
            states: Vec::with_capacity(count),
            f: Arc::new(f),
        };

        for _ in 0..count {
            cooks.add_cook();
        }

        cooks
    }

    fn add_cook(&mut self) {
        let (ttx, trx) = sync_channel::<Value<Option<Txty>>>(0);
        let (rtx, rrx) = sync_channel::<Value<Option<Rxty>>>(0);
        self.txs.push(ttx);
        self.rxs.push(rrx);
        self.states.push(Idling);
        let f = self.f.clone();
        Thread::spawn(move|| {
            loop {
                let input = trx.recv();
                if let Ok(Plain(Some(input))) = input {
                    if let Err(e) = rtx.send(Plain((*f)(input))) {
                        panic!("Slave could not send result ({})", e);
                    }
                } else if let Ok(Indexed(i, Some(input))) = input {
                    if let Err(e) = rtx.send(Indexed(i, (*f)(input))) {
                        panic!("Slave could not send result ({})", e);
                    }
                } else {
                    break;
                }
            }
        });
    }

    fn replace_cook(&mut self, i: usize) {
        let _ = self.txs.remove(i);
        let _ = self.rxs.remove(i);
        let _ = self.states.remove(i);
        self.add_cook();
    }

    pub fn cook<'a, T>(&'a mut self, it: T) -> Dishes<'a, Txty, Rxty, T, F>
    where T: Iterator<Item=Txty> {
        let count = self.txs.len();
        Dishes {
            cooks: self,
            it: it,
            count: count,
        }
    }

    pub fn cook_vec(&mut self, vec: Vec<Txty>) -> Vec<Option<Rxty>> {
        let mut output: Vec<Option<Rxty>> = Vec::with_capacity(vec.len());
        for _ in 0..vec.len() {
            output.push(None);
        }
        for (i, x) in vec.into_iter().enumerate() {
            if let Indexed(j, y) = self.send(Indexed(i, Some(x))) {
                output[j] = y;
            }
        }
        for i in 0..self.txs.len() {
            if self.states[i] == Working {
                if let Ok(Indexed(j, y)) = self.rxs[i].recv() {
                    output[j] = y;
                }
                self.states[i] = Idling;
            }
        }
        return output;
    }

    fn cook_owned<T>(self, it: T) -> OwnedDishes<Txty, Rxty, T, F>
    where T:Iterator<Item=Txty> {
        let count = self.txs.len();
        OwnedDishes {
            cooks: self,
            it: it,
            count: count,
        }
    }

    fn send(&mut self, x: Value<Option<Txty>>) -> Value<Option<Rxty>>  {
        let mut output: Value<Option<Rxty>> = Plain(None);
        let mut index: isize = -1;
        let mut broken = false;
        for i in 0..self.txs.len() {
            if self.states[i] == Idling {
                index = i as isize;
            }
        }
        if index < 0 {
            let select = Select::new();
            let mut hs = Vec::new();
            for rx in self.rxs.iter() {
                hs.push(select.handle(rx));
            }
            for h in hs.iter_mut() {
                unsafe{ h.add(); }
            }
            let id = select.wait();
            for (i, h) in hs.iter().enumerate() {
                if id == h.id() {
                    index = i as isize;
                    match self.rxs[i].recv() {
                        Ok(r) => {
                            output = r;
                        },
                        Err(_) => {
                            broken = true;
                        }
                    }
                }
            }
            for h in hs.iter_mut() {
                unsafe { h.remove(); }
            }
        }
        if broken {
            self.replace_cook(index as usize);
            return self.send(x);
        }
        if let Plain(None) = x {
            self.states[index as usize] = Done;
            return Plain(None);
        } else {
            self.states[index as usize] = Working;
        }
        match self.txs[index as usize].send(x) {
            Ok(_) => return output,
            Err(SendError(x)) => {
                if let Plain(None) = output {
                    self.replace_cook(index as usize);
                    return self.send(x);
                } else {
                    panic!("Master could not send data.");
                }
            }
        }
    }
}

#[unsafe_destructor]
impl<Txty: Send, Rxty: Send, F> Drop for Cooks<Txty, Rxty, F> {
    fn drop(&mut self) {
        for i in 0..self.txs.len() {
            if self.states[i] == Working {
                let _ = self.rxs[i].recv();
                let _ = self.txs[i].send(Plain(None));
            }
        }
    }
}

pub struct Dishes<'a, Txty: 'a, Rxty: 'a, T, F: 'a> {
    cooks: &'a mut Cooks<Txty, Rxty, F>,
    it: T,
    count: usize,
}

impl<'a, Txty: Send, Rxty: Send, T, F> Iterator for Dishes<'a, Txty, Rxty, T, F>
where T: Iterator<Item=Txty>,
      F: Fn(Txty) -> Option<Rxty>, F: Send + Sync + 'static {
    type Item = Rxty;

    fn next(&mut self) -> Option<Rxty> {
        loop {
            if let Some(input) = self.it.next() {
                if let Plain(Some(output)) = self.cooks.send(Plain(Some(input))) {
                    return Some(output);
                } else {
                    continue;
                }
            } else {
                if self.count > 0 {
                    self.count -= 1;
                    if self.cooks.states[self.count] == Working {
                        self.cooks.states[self.count] = Idling;
                        match self.cooks.rxs[self.count].recv() {
                            Err(_) => return self.next(),
                            Ok(Plain(output)) => return output,
                            Ok(_) => unreachable!(),
                        }
                    } else {
                        continue;
                    }
                } else {
                    return None;
                }
            }
        }
    }
}

pub struct OwnedDishes<Txty, Rxty, T, F> {
    cooks: Cooks<Txty, Rxty, F>,
    it: T,
    count: usize,
}

impl<Txty: Send, Rxty: Send, T, F> Iterator for OwnedDishes<Txty, Rxty, T, F>
where T: Iterator<Item=Txty>,
      F: Fn(Txty) -> Option<Rxty>, F: Send + Sync + 'static {
    type Item = Rxty;

    fn next(&mut self) -> Option<Rxty> {
        loop {
            if let Some(input) = self.it.next() {
                if let Plain(Some(output)) = self.cooks.send(Plain(Some(input))) {
                    return Some(output);
                } else {
                    continue;
                }
            } else {
                if self.count > 0 {
                    self.count -= 1;
                    if self.cooks.states[self.count] == Working {
                        if let Plain(output) = self.cooks.rxs[self.count].recv().unwrap() {
                            self.cooks.states[self.count] = Idling;
                            return output;
                        } else {
                            unreachable!();
                        }
                    } else {
                        continue;
                    }
                } else {
                    return None;
                }
            }
        }
    }
}

pub trait Cookable<Rxty, F> {
    type Input;
    type Output;

    /// Example:
    ///
    /// ```rust
    /// use cooks::Cookable;
    ///
    /// let vec = vec![1i32, 2, 3, 4];
    /// let result = vec![Some(1f32), Some(2f32), Some(3f32), Some(4f32)];
    ///
    /// assert_eq!(vec.cook(|&: x: i32| Some(x as f32)), result);
    /// ```
    ///
    /// Example with failing cooks:
    ///
    /// ```rust
    /// use cooks::Cookable;
    ///
    /// let vec = vec![1i32, 2, 3, 4];
    /// let result = vec![Some(1f32), Some(2f32), None, Some(4f32)];
    ///
    /// assert_eq!(vec.cook(|&: x: i32| {
    ///     if x == 3 { panic!("Does not cook") }
    ///     Some(x as f32)
    /// }), result);
    /// ```
    fn cook(self, f: F) -> Self::Output;

    fn cook_n(self, count: usize, f: F) -> Self::Output;
}

impl<Txty: Send, Rxty: Send, F> Cookable<Rxty, F> for Vec<Txty>
where F: Fn(Txty) -> Option<Rxty>, F: Send + Sync + 'static {
    type Input = Txty;
    type Output = Vec<Option<Rxty>>;

    fn cook(self, f: F) -> Vec<Option<Rxty>> {
        self.cook_n(::std::os::num_cpus(), f)
    }

    fn cook_n(self, count: usize, f: F) -> Vec<Option<Rxty>> {
        let mut output: Vec<Option<Rxty>> = Vec::with_capacity(self.len());
        for _ in 0..self.len() {
            output.push(None);
        }
        let mut cooks = Cooks::new_n(count, move|(i, x): (usize, Txty)| {
            Some((i, f(x)))
        });
        for (i, x) in cooks.cook(self.into_iter().enumerate()) {
            output[i] = x;
        }
        output
    }
}

impl<'a, Rxty: Send, Txty: Send, T, F> Cookable<Rxty, F> for T
where T: Iterator<Item=Txty>,
      F: Fn(Txty) -> Option<Rxty>, F: Send + Sync + 'static {
    type Input = Txty;
    type Output = OwnedDishes<Txty, Rxty, T, F>;

    fn cook(self, f: F) -> OwnedDishes<Txty, Rxty, T, F> {
        self.cook_n(::std::os::num_cpus(), f)
    }

    fn cook_n(self, count: usize, f: F) -> OwnedDishes<Txty, Rxty, T, F> {
        let cooks = Cooks::new_n(count, f);
        cooks.cook_owned(self)
    }
}

#[cfg(test)]
mod test {
    use std::iter::AdditiveIterator;
    use super::{Cooks, Cookable};

    #[test]
    fn test_cook() {
        for i in 0i32..20 {
            let mut cooks = Cooks::new_n((20 - i) as usize, |x: i32| {
                Some(x as f32)
            });
            let mut result = 0f32;
            for j in cooks.cook(0i32..(i + 1)) {
                result += j;
            }
            let mut result2 = 0f32;
            for j in cooks.cook(0i32..(i + 1)) {
                result2 += j;
            }
            assert_eq!((0i32..(i + 1)).sum() as f32, result);
            assert_eq!(result, result2);
        }
    }

    #[test]
    fn test_cook_vec() {
        for i in 0us..20 {
            let mut cooks = Cooks::new_n((20 - i), |x: i32| {
                Some(x as f32)
            });
            for (i, x) in cooks.cook_vec((0i32..((i+1) as i32)).collect::<Vec<i32>>())
                               .iter().enumerate() {
                assert_eq!(Some(i as f32), *x);
            }
            for (i, x) in cooks.cook_vec((0i32..((i+1) as i32)).collect::<Vec<i32>>())
                               .iter().enumerate() {
                assert_eq!(Some(i as f32), *x);
            }
        }
    }

    #[test]
    fn test_failing_slave() {

        let mut cooks = Cooks::new_n(8, |x: i32| {
            if x != 10 {
                Some(x as f32)
            } else {
                panic!("PANIC!!1");
            }
        });

        let mut result = 0f32;
        for x in cooks.cook(0i32..20) {
            result += x;
        }
        assert_eq!(((0i32..20).sum() - 10) as f32, result);

        let mut result = 0f32;
        for x in cooks.cook(0i32..20) {
            result += x;
        }
        assert_eq!(((0i32..20).sum() - 10) as f32, result);
    }

    #[test]
    fn test_cookable() {

        for i in 0i32..20 {
            for (i, x) in (0i32..(i+1)).collect::<Vec<i32>>()
                                       .cook_n((20 - i) as usize, |&: x: i32| {
                                            Some(x as f32)
                                       }).iter().enumerate() {
                assert_eq!(Some(i as f32), *x);
            }

            let mut result = 0f32;
            for j in (0i32..(i + 1)).cook_n((20 - i) as usize, |&: x: i32| Some(x as f32)) {
                result += j;
            }
            assert_eq!((0i32..(i + 1)).sum() as f32, result);
        }
    }
}