use std::ops::Add;
use std::ops::Mul;

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::tag::Tag;

#[derive(BlockIO)]
pub struct Adder<Ia, Ib, O>
where
    Ia: Add<Ib, Output = O> + 'static,
    Ib: 'static,
    O: 'static,
{
    #[input]
    input_a: In<Ia>,

    #[input]
    input_b: In<Ib>,

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

impl<Ia, Ib, O> Adder<Ia, Ib, O>
where
    Ia: Add<Ib, Output = O> + 'static,
    Ib: 'static,
    O: 'static,
{
    pub fn new(input_a: In<Ia>, input_b: In<Ib>) -> (Self, In<O>)
    {
        let (output, added) = oxydsp_flowgraph::io::stream();
        (
            Self {
                input_a,
                input_b,
                output,
            },
            added,
        )
    }
}

impl<Ia, Ib, O> Block for Adder<Ia, Ib, O>
where
    Ia: Add<Ib, Output = O> + 'static,
    Ib: 'static,
    O: 'static,
{
    fn work(&mut self) -> BlockResult
    {
        self.output.push_iter(
            self.input_a
                .iter()
                .zip(self.input_b.iter())
                .map(|(x, y)| {
                    (x.0 + y.0, Tag::from_tag_opts(&[&x.1, &y.1])).into()
                })
        );
        BlockResult::Ok
    }
}

#[derive(BlockIO)]
pub struct Multiplier<Ia, Ib, O>
where
    Ia: Mul<Ib, Output = O> + 'static,
    Ib: 'static,
    O: 'static,
{
    #[input]
    input_a: In<Ia>,

    #[input]
    input_b: In<Ib>,

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

impl<Ia, Ib, O> Multiplier<Ia, Ib, O>
where
    Ia: Mul<Ib, Output = O> + 'static,
    Ib: 'static,
    O: 'static,
{
    pub fn new(input_a: In<Ia>, input_b: In<Ib>) -> (Self, In<O>)
    {
        let (output, added) = oxydsp_flowgraph::io::stream();
        (
            Self {
                input_a,
                input_b,
                output,
            },
            added,
        )
    }
}

impl<Ia, Ib, O> Block for Multiplier<Ia, Ib, O>
where
    Ia: Mul<Ib, Output = O> + 'static,
    Ib: 'static,
    O: 'static,
{
    fn work(&mut self) -> BlockResult
    {
        self.output.push_iter(
            self.input_a
                .iter()
                .zip(self.input_b.iter())
                .map(|(x, y)| {
                    (x.0 * y.0, Tag::from_tag_opts(&[&x.1, &y.1])).into()
                })
        );
        BlockResult::Ok
    }
}
