distributed_physics_test/src/net/distribution.rs

use std::marker::PhantomData;

use bevy::{
    ecs::{component::Mutable, query::QueryFilter},
    prelude::*,
};
use uuid::Uuid;

use super::{
    packet::{InboundPacket, OutboundPacket, Packet},
    peer::PeerID,
};

/// Entities wishing to be networked must have this in their bundle
#[derive(Component, Debug)]
pub struct EntityNetworkID(Uuid);

impl Default for EntityNetworkID {
    fn default() -> Self {
        Self(Uuid::new_v4())
    }
}

#[derive(Component, Debug)]
pub struct PeerOwned;

pub trait NetworkEncodable {
    fn encode(&self) -> Vec<u8>;
}

impl<T> NetworkEncodable for T
where
    T: Clone + Into<Vec<u8>>,
{
    fn encode(&self) -> Vec<u8> {
        self.clone().into()
    }
}

pub trait NetworkDecodable: Sized {
    type DecodeError;
    fn decode(buffer: Vec<u8>) -> std::result::Result<Self, Self::DecodeError>;
}

impl<T> NetworkDecodable for T
where
    T: TryFrom<Vec<u8>>,
{
    type DecodeError = T::Error;
    fn decode(buffer: Vec<u8>) -> std::result::Result<Self, Self::DecodeError> {
        T::try_from(buffer)
    }
}

/// Components wishing to be networked must implement this type
pub trait Networked: Component<Mutability = Mutable> + NetworkEncodable + NetworkDecodable {}

impl<T> Networked for T where
    T: Component<Mutability = Mutable> + NetworkEncodable + NetworkDecodable
{
}

pub trait NetworkSettings: 'static + Send + Sync {
    type IncomingFilter: QueryFilter;
    type LocalChangeFilter: QueryFilter;
    type NewPeerFilter: QueryFilter;
    type NewLocalEntityFilter: QueryFilter;
}

pub struct DefaultNetworkSettings;

impl NetworkSettings for DefaultNetworkSettings {
    type IncomingFilter = With<PeerOwned>;
    type LocalChangeFilter = Without<PeerOwned>;
    type NewPeerFilter = Without<PeerOwned>;
    type NewLocalEntityFilter = Without<PeerOwned>;
}

fn incoming_network_entity<
    T: NetworkDecodable + Component<Mutability = Mutable>,
    F: QueryFilter,
>(
    mut inbound: MessageReader<InboundPacket>,
    mut components: Query<(&mut T, &EntityNetworkID), F>,
    mut commands: Commands,
) {
    'packets: for InboundPacket(packet) in inbound.read() {
        if let Ok(component) = T::decode(packet.message.clone()) {
            for (mut existing_component, id) in &mut components {
                if id.0 == packet.entity {
                    *existing_component = component;
                    continue 'packets;
                }
            }
            commands.spawn((component, EntityNetworkID(packet.entity), PeerOwned));
        }
    }
}

fn changed_local_entity<T: NetworkEncodable + Component, F: QueryFilter>(
    components: Query<(&T, &EntityNetworkID), (Changed<T>, F)>,
    peers: Query<&PeerID>,
    mut outbound: MessageWriter<OutboundPacket>,
) {
    for (component, id) in components {
        for peer in peers {
            outbound.write(Packet::create(peer.id, id.0, component.encode()));
        }
    }
}

fn new_peer<T: NetworkEncodable + Component, F: QueryFilter>(
    add: On<Add, PeerID>,
    components: Query<(&T, &EntityNetworkID), F>,
    peers: Query<&PeerID>,
    mut outbound: MessageWriter<OutboundPacket>,
) -> Result {
    let peer = peers.get(add.entity)?;
    for (component, id) in components {
        outbound.write(Packet::create(peer.id, id.0, component.encode()));
    }
    Ok(())
}

fn new_local_entity<T: NetworkEncodable + Component, F: QueryFilter>(
    add: On<Add, T>,
    components: Query<(&T, &EntityNetworkID), F>,
    peers: Query<&PeerID>,
    mut outbound: MessageWriter<OutboundPacket>,
) {
    if let Ok((component, id)) = components.get(add.entity) {
        for peer in peers {
            outbound.write(Packet::create(peer.id, id.0, component.encode()));
        }
    }
}

#[derive(Debug)]
pub struct DistributionPlugin<T: Networked, Settings: NetworkSettings = DefaultNetworkSettings> {
    _phantom1: PhantomData<T>,
    _phantom2: PhantomData<Settings>,
}

impl<T: Networked, Settings: NetworkSettings> Default for DistributionPlugin<T, Settings> {
    fn default() -> Self {
        Self {
            _phantom1: Default::default(),
            _phantom2: Default::default(),
        }
    }
}

impl<T: Networked, Settings: NetworkSettings> Plugin for DistributionPlugin<T, Settings> {
    fn build(&self, app: &mut App) {
        app.add_systems(
            FixedUpdate,
            (
                incoming_network_entity::<T, Settings::IncomingFilter>,
                changed_local_entity::<T, Settings::LocalChangeFilter>,
            ),
        )
        .add_observer(new_peer::<T, Settings::NewPeerFilter>)
        .add_observer(new_local_entity::<T, Settings::NewLocalEntityFilter>);
    }
}