#![allow(clippy::all)]

use crate::*;
use bevy_core::Name;
use bevy_ecs::prelude::*;
use criterion::*;
use std::fmt::Display;

/// Define a trait for our components to implement.
#[queryable]
pub trait Messages {
    fn messages(&self) -> &[String];
    fn send_message(&mut self, _: &dyn Display);
}

#[derive(Component)]
pub struct RecA {
    messages: Vec<String>,
}

impl Messages for RecA {
    fn messages(&self) -> &[String] {
        &self.messages
    }
    fn send_message(&mut self, msg: &dyn Display) {
        self.messages.push(msg.to_string());
    }
}

#[derive(Component)]
pub struct RecB {
    messages: Vec<String>,
}

impl Messages for RecB {
    fn messages(&self) -> &[String] {
        &self.messages
    }
    fn send_message(&mut self, msg: &dyn Display) {
        self.messages.push(msg.to_string());
    }
}

pub struct Benchmark<'w>(World, QueryState<&'w dyn Messages>, Vec<usize>);

impl<'w> Benchmark<'w> {
    // Each entity only has one component in practice.
    fn one() -> Self {
        let mut world = World::new();

        world.register_component_as::<dyn Messages, RecA>();
        world.register_component_as::<dyn Messages, RecB>();

        for _ in 0..5_000 {
            world.spawn((Name::new("Hello"), RecA { messages: vec![] }));
        }
        for _ in 0..5_000 {
            world.spawn((Name::new("Hello"), RecB { messages: vec![] }));
        }

        let query = world.query();
        Self(world, query, Default::default())
    }
    fn multiple() -> Self {
        let mut world = World::new();

        world.register_component_as::<dyn Messages, RecA>();
        world.register_component_as::<dyn Messages, RecB>();

        for _ in 0..10_000 {
            world.spawn((
                Name::new("Hello"),
                RecA { messages: vec![] },
                RecB { messages: vec![] },
            ));
        }

        let query = world.query();
        Self(world, query, Default::default())
    }
    // Queries with only one, and queries with multiple.
    pub fn distributed() -> Self {
        let mut world = World::new();

        world.register_component_as::<dyn Messages, RecA>();
        world.register_component_as::<dyn Messages, RecB>();

        for _ in 0..2_500 {
            world.spawn((Name::new("Hello"), RecA { messages: vec![] }));
        }
        for _ in 0..2_500 {
            world.spawn((Name::new("Hello"), RecB { messages: vec![] }));
        }
        for _ in 0..5_000 {
            world.spawn((
                Name::new("Hello"),
                RecA { messages: vec![] },
                RecB { messages: vec![] },
            ));
        }

        let query = world.query();
        Self(world, query, Default::default())
    }

    pub fn run(&mut self) {
        let mut output = Vec::new();
        for all in self.1.iter_mut(&mut self.0) {
            for x in all {
                output.push(x.messages().len());
            }
        }
        self.2 = output;
    }
}

pub fn one(c: &mut Criterion) {
    let mut benchmark = Benchmark::one();
    c.bench_function("All<> - 1 match", |b| b.iter(|| benchmark.run()));
    eprintln!("{}", benchmark.2.len());
}
pub fn multiple(c: &mut Criterion) {
    let mut benchmark = Benchmark::multiple();
    c.bench_function("All<> - 2 matches", |b| b.iter(|| benchmark.run()));
    eprintln!("{}", benchmark.2.len());
}
pub fn distributed(c: &mut Criterion) {
    let mut benchmark = Benchmark::distributed();
    c.bench_function("All<> - 1-2 matches", |b| b.iter(|| benchmark.run()));
    eprintln!("{}", benchmark.2.len());
}

criterion_group!(all, one, multiple, distributed);
criterion_main!(all);