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;
use aeronet_websocket::tungstenite::Bytes;
use avian2d::prelude::{AngularInertia, AngularVelocity, Collider, LinearVelocity, Mass, Position, Rotation};
use bevy::ecs::entity::MapEntities;
use bevy::ecs::system::SystemState;
use bevy::prelude::*;
use postcard::{from_bytes, to_allocvec, to_slice, to_vec};
use crate::prelude::{App, Message};
use crate::shared::thrust::ThrustSolution;
//use bevy_replicon::prelude::*;
use serde::{Deserialize, Serialize};
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};


/*pub fn register_replication(app: &mut App) {
    app
        .add_mapped_server_message::<Hi>(Channel::Ordered)

        .replicate::<Transform>()
        .replicate::<GlobalTransform>()

        .replicate::<Position>()
        .replicate::<Rotation>()
        .replicate::<LinearVelocity>()
        .replicate::<AngularVelocity>()
        .replicate::<AngularInertia>()
        .replicate::<Mass>()
        .replicate::<Collider>()

        .replicate::<Part>()
        .replicate::<Planet>()
        .replicate::<Player>()

        .replicate::<ChildOf>()
        .replicate::<Ship>()
        .replicate::<PartInShip>()
        .replicate::<Joint>()
        .replicate::<Peer>()
        .replicate::<JointOf>()
        .replicate::<SnapOfJoint>()
        .replicate::<SnapOf>()
        .replicate::<PlayerStorage>()
        .replicate::<Thruster>()
        .replicate::<ThrusterOfPart>()
        .replicate::<CanCraft>()
        .replicate::<Temperature>()
        .replicate::<Drill>()
        .replicate::<SingleStorage>()
        .replicate::<PlanetSpring>()
        .replicate::<PlanetSpringJoint>();
}*/

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

        .add_client_message::<DragRequestEvent>()
        .add_client_message::<ToggleDrillEvent>()
        .add_client_message::<CraftPartRequest>()
        .add_client_message::<ThrustSolution>();
}

#[derive(Message, 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.add_systems(PostUpdate, recv_from_server);
    app.add_systems(PostUpdate, 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 ServerMessageRegistry {
    forward: HashMap<String, LaneIndex>,
    reverse: HashMap<LaneIndex, fn(Vec<u8>, &mut World)>,
}
#[derive(Resource, Default)]
pub struct ClientMessageRegistry {
    forward: HashMap<String, 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_server_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self) -> &mut Self;
    fn add_client_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self) -> &mut Self;
}
impl NetAppExt for App {
    fn add_server_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self) -> &mut Self {
        self
            .add_message::<ToClients<T>>()
            .add_message::<T>()
            .add_systems(Startup, register_server_message::<T>)
            .add_systems(Update, send_to_client::<T>)
    }
    fn add_client_message<T: Message + TypePath + Serialize + for<'a> Deserialize<'a>>(&mut self) -> &mut Self {
        self
            .add_message::<FromClients<T>>()
            .add_message::<T>()
            .add_systems(Startup, register_client_message::<T>)
            .add_systems(Update, send_to_server::<T>)
    }
}
fn register_server_message<T: Message + TypePath + for<'a> Deserialize<'a>>(mut registry: ResMut<ServerMessageRegistry>) {
    registry.forward.insert(T::type_path().to_string(), get_lane_index());
    registry.reverse.insert(get_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>>(mut registry: ResMut<ClientMessageRegistry>) {
    registry.forward.insert(T::type_path().to_string(), get_lane_index());
    registry.reverse.insert(get_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(M::type_path())
                            .expect("Failed to get message lane; the message likely isn't serialized yet"),
                        to_allocvec(&message).expect("Failed to serialize message").into(),
                        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(M::type_path())
                                .expect("Failed to get message lane; the message likely isn't serialized yet"),
                            to_allocvec(&message).expect("Failed to serialize message").into(),
                            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(M::type_path())
                    .expect("Failed to get message lane; the message likely isn't serialized yet"),
                to_allocvec(&message).expect("Failed to serialize message").into(),
                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 = Vec::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 = Vec::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);
    }
}