pub mod planet;
pub mod staged_transform;
pub mod part;
pub mod was_updated;

use std::any::TypeId;
use std::collections::HashMap;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
use std::time::Instant;

use aeronet_transport::Transport;
use aeronet_transport::lane::{LaneIndex, LaneKind};
use aeronet_websocket::tungstenite::Bytes;
use avian2d::prelude::{AngularInertia, AngularVelocity, Collider, LinearVelocity, Mass, Position, Rotation};
use bevy::ecs::entity::{EntityHashMap, MapEntities};
use bevy::ecs::system::SystemState;
use bevy::prelude::*;
use bevy::reflect::DynamicTypePath;
use postcard::{from_bytes, to_allocvec, to_slice, to_vec};
use crate::prelude::{App, Message};
use crate::shared::thrust::ThrustSolution;
use serde::{Deserialize, Serialize};
use smolvec::SmolVec;
use crate::shared::attachment::{Joint, JointOf, PartInShip, Peer, Ship, SnapOf, SnapOfJoint};
use crate::shared::config::planet::{Planet, PlanetSpring, PlanetSpringJoint};
use crate::shared::ecs::{CanCraft, CraftPartRequest, DragRequestEvent, Drill, Part, Player, PlayerStorage, SingleStorage, Temperature, ToggleDrillEvent};
use crate::shared::ecs::thruster::{Thruster, ThrusterOfPart};
use crate::shared::net::part::PartDto;
use crate::shared::net::planet::PlanetDto;

pub fn register_net(app: &mut App) {
    app
        .add_server_message::<Hi>(LaneKind::ReliableOrdered)

        .add_server_message::<PlanetDto>(LaneKind::UnreliableSequenced)
        .add_server_message::<PartDto>(LaneKind::UnreliableSequenced)

        .add_client_message::<DragRequestEvent>(LaneKind::ReliableOrdered)
        .add_client_message::<ToggleDrillEvent>(LaneKind::ReliableOrdered)
        .add_client_message::<CraftPartRequest>(LaneKind::ReliableOrdered)
        .add_client_message::<ThrustSolution>(LaneKind::ReliableOrdered)

        .add_systems(PostStartup, |srv_msg_reg: Res<ServerMessageRegistry>, client_msg_reg: Res<ClientMessageRegistry>| {
            debug!("network registration complete");
            debug!("client registry: {:#?}", client_msg_reg);
            debug!("server registry: {:#?}", srv_msg_reg);
        });
}

#[derive(Message, Clone, Deserialize, Serialize, TypePath, MapEntities)]
pub struct Hi {
    #[entities]
    pub you_are: Entity,
    pub time_offset: f64,
}

pub fn setup_net(app: &mut App) {
    app.insert_resource(ServerMessageRegistry::default());
    app.insert_resource(ClientMessageRegistry::default());
    app.insert_resource(TransportLanes::default());
    app.add_systems(Update, recv_from_server);
    app.add_systems(Update, recv_from_client);
}
#[derive(Message, Deref, Deserialize, Serialize)]
pub struct ToClients<T: Message + TypePath> {
    #[deref]
    pub message: T,
    pub targets: SendTargets,
}
#[derive(Message, Deref, Deserialize, Serialize)]
pub struct FromClients<T: Message + TypePath> {
    pub client_id: ClientId,
    #[deref]
    pub message: T,
}

#[derive(Deserialize, Serialize)]
pub enum SendTargets {
    All,
    Single(ClientId),
}

#[derive(Deserialize, Serialize)]
pub enum ClientId {
    Server,
    Client(Entity),
}

#[derive(Resource, Default)]
pub struct TransportLanes {
    pub lanes: Vec<LaneKind>
}

#[derive(Resource, Default)]
pub struct ServerEntityMap {
    pub server_to_client: EntityHashMap<Entity>,
    pub client_to_server: EntityHashMap<Entity>,
}
#[derive(Resource, Default, Debug)]
pub struct ServerMessageRegistry {
    forward: HashMap<TypeId, LaneIndex>,
    reverse: HashMap<LaneIndex, fn(Vec<u8>, &mut World)>,
}
#[derive(Resource, Default, Debug)]
pub struct ClientMessageRegistry {
    forward: HashMap<TypeId, LaneIndex>,
    reverse: HashMap<LaneIndex, fn(Vec<u8>, ClientId, &mut World)>,
}

static COUNTER: AtomicU32 = AtomicU32::new(0);
fn get_lane_index() -> LaneIndex {
    LaneIndex::new(COUNTER.fetch_add(1, Ordering::Relaxed))
}

pub trait NetAppExt {
    fn add_mapped_server_message<T: Message + MapEntities + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self;
    fn add_server_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self;
    fn add_mapped_client_message<T: Message + Clone + MapEntities + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self;
    fn add_client_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self;
}
impl NetAppExt for App {
    fn add_mapped_server_message<T: Message + MapEntities + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self {
        let mut registry = self.world_mut().resource_mut::<ServerMessageRegistry>();
        register_mapped_server_message::<T>(&mut registry);
        self.world_mut().resource_mut::<TransportLanes>().lanes.push(kind);
        self
            .add_message::<ToClients<T>>()
            .add_message::<T>()
            .add_systems(Update, send_to_client::<T>)
    }
    fn add_server_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self {
        let mut registry = self.world_mut().resource_mut::<ServerMessageRegistry>();
        register_server_message::<T>(&mut registry);
        self.world_mut().resource_mut::<TransportLanes>().lanes.push(kind);
        self
            .add_message::<ToClients<T>>()
            .add_message::<T>()
            .add_systems(Update, send_to_client::<T>)
    }
    fn add_mapped_client_message<T: Message + Clone + MapEntities + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self {
        let mut registry = self.world_mut().resource_mut::<ClientMessageRegistry>();
        register_client_message::<T>(&mut registry);
        self.world_mut().resource_mut::<TransportLanes>().lanes.push(kind);
        self
            .add_message::<FromClients<T>>()
            .add_message::<T>()
            .add_systems(Update, send_mapped_to_server::<T>)
    }
    fn add_client_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self, kind: LaneKind) -> &mut Self {
        let mut registry = self.world_mut().resource_mut::<ClientMessageRegistry>();
        register_client_message::<T>(&mut registry);
        self.world_mut().resource_mut::<TransportLanes>().lanes.push(kind);
        self
            .add_message::<FromClients<T>>()
            .add_message::<T>()
            .add_systems(Update, send_to_server::<T>)
    }
}
fn register_mapped_server_message<T: Message + MapEntities + TypePath + for<'a> Deserialize<'a>>(registry: &mut ServerMessageRegistry) {
    let lane_index = get_lane_index();
    registry.forward.insert(TypeId::of::<T>(), lane_index);
    registry.reverse.insert(lane_index, |payload: Vec<u8>, world: &mut World| {
        let mut entity_map = world.resource_mut::<ServerEntityMap>();
        let mut message = from_bytes::<T>(&payload).expect(&format!("Failed to deserialize message of type {}", T::type_path()));
        message.map_entities(&mut entity_map.server_to_client);
        world.write_message(message).expect("Could not send message to game");
    });
}
fn register_server_message<T: Message + TypePath + for<'a> Deserialize<'a>>(registry: &mut ServerMessageRegistry) {
    let lane_index = get_lane_index();
    registry.forward.insert(TypeId::of::<T>(), lane_index);
    registry.reverse.insert(lane_index, |payload: Vec<u8>, world: &mut World| {
        let message = from_bytes::<T>(&payload).expect(&format!("Failed to deserialize message of type {}", T::type_path()));
        world.write_message(message).expect("Could not send message to game");
    });
}
fn register_client_message<T: Message + TypePath + for<'a> Deserialize<'a>>(registry: &mut ClientMessageRegistry) {
    let lane_index = get_lane_index();
    registry.forward.insert(TypeId::of::<T>(), lane_index);
    registry.reverse.insert(lane_index, |payload: Vec<u8>, client_id: ClientId, world: &mut World| {
        let message = from_bytes::<T>(&payload).expect(&format!("Failed to deserialize message of type {}", T::type_path()));
        world.write_message(FromClients {
            client_id,
            message,
        }).expect("Could not send message to game");
    });
}

fn send_to_client<M: Message + TypePath + Serialize>(
    mut messages: MessageReader<ToClients<M>>,
    mut clients: Query<(Entity, &mut Transport), With<ChildOf>>,
    message_registry: Res<ServerMessageRegistry>,
) {
    for message in messages.read() {
        match &message.targets {
            SendTargets::All => {
                for (client_entity, mut transport) in &mut clients {
                    transport.send.push(
                        *message_registry.forward.get(&TypeId::of::<M>())
                            .expect("Failed to get message lane; the message likely isn't serialized yet"),
                        to_allocvec(&message).expect("Failed to serialize message").into(),
                        bevy::platform::time::Instant::now(),
                    ).expect("Failed to send message");
                }
            }
            SendTargets::Single(client) => {
                for (client_entity, mut transport) in &mut clients {
                    if let ClientId::Client(client) = client && client_entity == *client {
                        transport.send.push(
                            *message_registry.forward.get(&TypeId::of::<M>())
                                .expect("Failed to get message lane; the message likely isn't serialized yet"),
                            to_allocvec(&message).expect("Failed to serialize message").into(),
                            bevy::platform::time::Instant::now(),
                        ).expect("Failed to send message");
                    }
                }
            }
        }
    }
}
fn send_to_server<M: Message + TypePath + Serialize>(
    mut messages: MessageReader<M>,
    mut sessions: Query<&mut Transport, Without<ChildOf>>,
    message_registry: Res<ClientMessageRegistry>,
) {
    for message in messages.read() {
        for mut transport in &mut sessions {
            transport.send.push(
                *message_registry.forward.get(&TypeId::of::<M>())
                    .expect("Failed to get message lane; the message likely isn't serialized yet"),
                to_allocvec(&message).expect("Failed to serialize message").into(),
                bevy::platform::time::Instant::now(),
            ).expect("Failed to send message");
        }
    }
}
fn send_mapped_to_server<M: Message + Clone + MapEntities + TypePath + Serialize>(
    mut messages: MessageReader<M>,
    mut sessions: Query<&mut Transport, Without<ChildOf>>,
    mut entity_map: ResMut<ServerEntityMap>,
    message_registry: Res<ClientMessageRegistry>,
) {
    for message in messages.read() {
        for mut transport in &mut sessions {
            let mut message = message.clone();
            message.map_entities(&mut entity_map.client_to_server);
            transport.send.push(
                *message_registry.forward.get(&TypeId::of::<M>())
                    .expect("Failed to get message lane; the message likely isn't serialized yet"),
                to_allocvec(&message).expect("Failed to serialize message").into(),
                bevy::platform::time::Instant::now(),
            ).expect("Failed to send message");
        }
    }
}
fn recv_from_server(
    world: &mut World
) {
    let mut system_state: SystemState<(
        Query<&mut Transport, Without<ChildOf>>,
        Res<ServerMessageRegistry>,
    )> = SystemState::new(world);
    let (mut sessions, message_registry) = system_state.get_mut(world);
    let mut messages = SmolVec::new();
    for mut transport in sessions.iter_mut() {
        for message in transport.recv.msgs.drain() {
            let payload = message.payload;
            let message_fn = message_registry.reverse.get(&message.lane).expect("Packet was sent across a lane that didn't have a message assigned to it yet");
            messages.push((*message_fn, payload));
        }
        for _ in transport.recv.acks.drain() {

        }
    }
    for (message_fn, payload) in messages {
        message_fn(payload, world);
    }
}
fn recv_from_client(
    world: &mut World
) {
    let mut system_state: SystemState<(
        Query<(Entity, &mut Transport), With<ChildOf>>,
        Res<ClientMessageRegistry>,
    )> = SystemState::new(world);
    let (mut clients, message_registry) = system_state.get_mut(world);
    let mut messages = SmolVec::new();
    for (client_entity, mut transport) in clients.iter_mut() {
        for message in transport.recv.msgs.drain() {
            let payload = message.payload;
            let message_fn = message_registry.reverse.get(&message.lane).expect("Packet was sent across a lane that didn't have a message assigned to it yet");
            messages.push((payload, ClientId::Client(client_entity), *message_fn));
        }
        for _ in transport.recv.acks.drain() {

        }
    }
    for (payload, client_id, message_fn) in messages {
        message_fn(payload, client_id, world);
    }
}