use std::iter::FusedIterator;

use oxydsp_flowgraph::BlockIO;
use oxydsp_flowgraph::block::Block;
use oxydsp_flowgraph::block::BlockResult;
use oxydsp_flowgraph::io::In;
use oxydsp_flowgraph::io::Out;
use oxydsp_flowgraph::io::stream;
use oxydsp_flowgraph::tag::Tag;
use oxydsp_flowgraph::tag::TagMergable;
use oxydsp_flowgraph::tag::Tagged;

#[derive(BlockIO)]
pub struct Map<I: 'static, O: 'static, F>
{
    #[input]
    input: In<I>,

    #[output]
    output: Out<O>,

    map: F,
}

impl<I: 'static, O: 'static, F> Map<I, O, F>
where
    F: Fn(I) -> O,
{
    pub fn new(input: In<I>, map: F) -> (Self, In<O>)
    {
        let (output, mapped) = stream();
        (Self { input, output, map }, mapped)
    }
}

impl<I: 'static, O: 'static, F> Block for Map<I, O, F>
where
    F: Fn(I) -> O,
{
    fn work(&mut self) -> oxydsp_flowgraph::block::BlockResult
    {
        self.output.push_iter(
            self.input
                .iter()
                .map(|x| ((&self.map)(x.0), x.1).into()),
        );
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct MapResult<I: 'static, O: 'static, F>
{
    #[input]
    input: In<I>,

    #[output]
    output: Out<O>,

    map: F,
}

impl<I: 'static, O: 'static, F> MapResult<I, O, F>
where
    F: Fn(I) -> (O, BlockResult),
{
    pub fn new(input: In<I>, map: F) -> (Self, In<O>)
    {
        let (output, mapped) = stream();
        (Self { input, output, map }, mapped)
    }
}

impl<I: 'static, O: 'static, F> Block for MapResult<I, O, F>
where
    F: Fn(I) -> (O, BlockResult),
{
    fn work(&mut self) -> oxydsp_flowgraph::block::BlockResult
    {
        let mut writer = self.output.write_push();
        let mut reader = self.input.iter();

        for _ in 0..(writer.len().min(reader.len()))
        {
            let (input, tag_opt) = reader.next().unwrap().into();
            let (output, result) = (self.map)(input);
            let _ = writer.push((output, tag_opt).into());
            match result
            {
                BlockResult::Terminated | BlockResult::Exit =>
                {
                    return result;
                }
                BlockResult::Ok =>
                {}
            }
        }
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct MapResultTagged<I: 'static, O: 'static, F>
{
    #[input]
    input: In<I>,

    #[output]
    output: Out<O>,

    map: F,
}

impl<I: 'static, O: 'static, F> MapResultTagged<I, O, F>
where
    F: Fn(Tagged<I>) -> (Tagged<O>, BlockResult),
{
    pub fn new(input: In<I>, map: F) -> (Self, In<O>)
    {
        let (output, mapped) = stream();
        (Self { input, output, map }, mapped)
    }
}

impl<I: 'static, O: 'static, F> Block for MapResultTagged<I, O, F>
where
    F: Fn(Tagged<I>) -> (Tagged<O>, BlockResult),
{
    fn work(&mut self) -> oxydsp_flowgraph::block::BlockResult
    {
        let mut writer = self.output.write_push();
        let mut reader = self.input.iter();

        for _ in 0..(writer.len().min(reader.len()))
        {
            let (input, tag_opt) = reader.next().unwrap().into();
            let (tagged_out, result) = (self.map)((input, tag_opt.clone()).into());
            let (output, tag_out) = tagged_out.into();

            let _ = writer.push((output, tag_opt.merge(&tag_out)).into());
            match result
            {
                BlockResult::Terminated | BlockResult::Exit =>
                {
                    return result;
                }
                BlockResult::Ok =>
                {}
            }
        }
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct Scan<I: 'static, O: 'static, S, F>
where
    F: Fn(&mut S, I) -> O,
{
    #[input]
    input: In<I>,

    #[output]
    output: Out<O>,

    state: S,

    map: F,
}

impl<I: 'static, O: 'static, S, F> Scan<I, O, S, F>
where
    F: Fn(&mut S, I) -> O,
{
    pub fn new(input: In<I>, initial_state: S, map: F) -> (Self, In<O>)
    {
        let (output, mapped) = stream();
        (
            Self {
                input,
                output,
                state: initial_state,
                map,
            },
            mapped,
        )
    }
}

impl<I, O, S, F> Block for Scan<I, O, S, F>
where
    I: 'static,
    O: 'static,
    F: Fn(&mut S, I) -> O,
{
    fn work(&mut self) -> BlockResult
    {
        self.output.push_iter(
            self.input
                .iter()
                .map(|x| ((self.map)(&mut self.state, x.0), x.1).into()),
        );
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct ScanTagged<I: 'static, O: 'static, S, F>
where
    F: Fn(&mut S, Tagged<I>) -> Tagged<O>,
{
    #[input]
    input: In<I>,

    #[output]
    output: Out<O>,

    state: S,

    map: F,
}

impl<I: 'static, O: 'static, S, F> ScanTagged<I, O, S, F>
where
    F: Fn(&mut S, Tagged<I>) -> Tagged<O>,
{
    pub fn new(input: In<I>, initial_state: S, map: F) -> (Self, In<O>)
    {
        let (output, mapped) = stream();
        (
            Self {
                input,
                output,
                state: initial_state,
                map,
            },
            mapped,
        )
    }
}

impl<I, O, S, F> Block for ScanTagged<I, O, S, F>
where
    I: 'static,
    O: 'static,
    F: Fn(&mut S, Tagged<I>) -> Tagged<O>,
{
    fn work(&mut self) -> BlockResult
    {
        self.output.push_iter(self.input.iter().map(|tagged| {
            let cloned_tag = tagged.1.clone();
            let (output, tag) = (self.map)(&mut self.state, (tagged.0, cloned_tag).into()).into();
            (output, Tag::from_tag_opts(&[&tagged.1, &tag])).into()
        }));
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct Repeat<T: 'static>
{
    #[input]
    input: In<T>,

    repetitions: usize,
    remaining: usize,
    current: Option<(T, Option<Tag>)>,

    #[output]
    output: Out<T>,
}

impl<T: 'static> Repeat<T>
{
    pub fn new(input: In<T>, repetitions: usize) -> (Self, In<T>)
    {
        let (output, repeated) = stream();
        (
            Self {
                input,
                repetitions,
                remaining: repetitions,
                current: None,
                output,
            },
            repeated,
        )
    }
}

impl<T: Clone + 'static> Block for Repeat<T>
{
    fn work(&mut self) -> BlockResult
    {
        let mut writer = self.output.write_push();
        let mut reader = self.input.iter();
        let len = writer.len().min(reader.len() / self.repetitions + 1);

        for _ in 0..len
        {
            if self.remaining == 0 || self.current.is_none()
            {
                if let Some(x) = reader.next()
                {
                    self.current = Some(x.into());
                    self.remaining = self.repetitions;
                }
                else
                {
                    return BlockResult::Ok;
                }
            }

            writer
                .push(self.current.clone().unwrap().into())
                .unwrap_or_else(|_| panic!());

            if let Some((_, tag)) = &mut self.current
            {
                *tag = None;
            }

            self.remaining -= 1;
        }

        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct NullSink<T: 'static>
{
    #[input]
    input: In<T>,
}

impl<T: 'static> NullSink<T>
{
    pub fn new(input: In<T>) -> Self
    {
        Self { input }
    }
}

impl<I: 'static> Block for NullSink<I>
{
    fn work(&mut self) -> BlockResult
    {
        self.input.iter().for_each(|_| ());
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct Tee<T: 'static + Clone>
{
    #[input]
    input: In<T>,

    #[output]
    output_a: Out<T>,

    #[output]
    output_b: Out<T>,
}

impl<T: 'static + Clone> Tee<T>
{
    pub fn new(input: In<T>) -> (Self, In<T>, In<T>)
    {
        let (output_a, port_a) = stream();
        let (output_b, port_b) = stream();
        (
            Self {
                input,
                output_a,
                output_b,
            },
            port_a,
            port_b,
        )
    }
}

impl<T: 'static + Clone> Block for Tee<T>
{
    fn work(&mut self) -> BlockResult
    {
        let mut writer_a = self.output_a.write_push();
        let mut writer_b = self.output_b.write_push();
        for x in self.input.iter().take(writer_a.len().min(writer_b.len()))
        {
            let _ = writer_a.push(x.clone());
            let _ = writer_b.push(x);
        }
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct FlatMap<I, O, F>
where
    I: 'static,
    O: IntoIterator + 'static,
    O::IntoIter: FusedIterator,
    F: Fn(I) -> O,
{
    #[input]
    input: In<I>,

    #[output]
    output: Out<O::Item>,

    current_iter: Option<O::IntoIter>,
    map: F,
}

impl<I, O, F> FlatMap<I, O, F>
where
    I: 'static,
    O: IntoIterator + 'static,
    O::IntoIter: FusedIterator,
    F: Fn(I) -> O,
{
    pub fn new(input: In<I>, map: F) -> (Self, In<O::Item>)
    {
        let (output, port) = oxydsp_flowgraph::io::stream();
        (
            Self {
                input,
                output,
                current_iter: None,
                map,
            },
            port,
        )
    }
}

impl<I, O, F> Block for FlatMap<I, O, F>
where
    I: 'static,
    O: IntoIterator + 'static,
    O::IntoIter: FusedIterator,
    F: Fn(I) -> O,
{
    fn work(&mut self) -> BlockResult
    {
        let mut writer = self.output.write_push();
        let mut reader = self.input.iter();

        let max_write = writer.len();
        let mut written = 0;

        while written < max_write
        {
            if let Some(current_iter) = self.current_iter.as_mut()
            {
                if let Some(next_elt) = current_iter.next()
                {
                    let _ = writer.push((next_elt, None).into());
                    written += 1;
                    continue;
                }
                else
                {
                    // Iterator empty
                    self.current_iter = None;
                }
            }

            if self.current_iter.is_none()
            {
                // Get input
                if let Some(input) = reader.next()
                {
                    let mut new_iter = (self.map)(input.0).into_iter();
                    if let Some(first_elt) = new_iter.next()
                    {
                        self.current_iter = Some(new_iter);
                        let _ = writer.push((first_elt, input.1).into());
                        written += 1;
                    }
                    else
                    {
                        // Iterator empty
                        self.current_iter = None;
                        continue;
                    }
                }
                else
                {
                    // Cannot continue
                    break;
                }
            }
        }

        BlockResult::Ok
    }
}
