using Unity.Burst;
using UnityEngine;
using Unity.Networking.Transport;
using Unity.Collections;
using Unity.Jobs;

namespace Unity.Networking.Transport.Samples
{
    public class CustomNetworkInterfaceConnect : MonoBehaviour
    {
        const int ServerPort = 9000;
        struct PendingPing
        {
            public int id;
            public float time;
        }

        private NetworkDriver m_ClientDriver;
        private NetworkConnection m_clientToServerConnection;
        public NetworkDriver m_ServerDriver;
        private NativeList<NetworkConnection> m_serverConnections;
        // pendingPing is a ping sent to the server which have not yet received a response.
        private PendingPing m_pendingPing;
        // The ping stats are two integers, time for last ping and number of pings
        private int m_lastPingTime;
        private int m_numPingsSent;

        void Start()
        {
            // Create a NetworkDriver for the client. We could bind to a specific address but in this case we rely on the
            // implicit bind since we do not need to bing to anything special
            m_ClientDriver = new NetworkDriver(new UDPNetworkInterface());

            // Create the server driver, bind it to a port and start listening for incoming connections
            m_ServerDriver = new NetworkDriver(new UDPNetworkInterface());
            var addr = NetworkEndPoint.AnyIpv4;
            addr.Port = ServerPort;
            if (m_ServerDriver.Bind(addr) != 0)
                Debug.Log($"Failed to bind to port {ServerPort}");
            else
                m_ServerDriver.Listen();

            m_serverConnections = new NativeList<NetworkConnection>(16, Allocator.Persistent);
        }

        void OnDestroy()
        {
            m_ClientDriver.Dispose();

            m_ServerDriver.Dispose();
            m_serverConnections.Dispose();
        }

        void FixedUpdate()
        {
            ClientUpdate();
            ServerUpdate();
        }

        void ClientUpdate()
        {
            // Update the NetworkDriver. It schedules a job so we must wait for that job with Complete
            m_ClientDriver.ScheduleUpdate().Complete();

            // If the client ui indicates we should be sending pings but we do not have an active connection we create one
            if (!m_clientToServerConnection.IsCreated)
            {
                var serverEP = NetworkEndPoint.LoopbackIpv4;
                serverEP.Port = ServerPort;
                m_clientToServerConnection = m_ClientDriver.Connect(serverEP);
            }

            DataStreamReader strm;
            NetworkEvent.Type cmd;
            // Process all events on the connection. If the connection is invalid it will return Empty immediately
            while ((cmd = m_clientToServerConnection.PopEvent(m_ClientDriver, out strm)) != NetworkEvent.Type.Empty)
            {
                if (cmd == NetworkEvent.Type.Connect)
                {
                    // When we get the connect message we can start sending data to the server
                    // Set the ping id to a sequence number for the new ping we are about to send
                    m_pendingPing = new PendingPing {id = m_numPingsSent, time = Time.fixedTime};
                    // Create a 4 byte data stream which we can store our ping sequence number in

                    if (m_ClientDriver.BeginSend(m_clientToServerConnection, out var pingData) == 0)
                    {
                        pingData.WriteInt(m_numPingsSent);
                        m_ClientDriver.EndSend(pingData);
                    }
                    // Update the number of sent pings
                    ++m_numPingsSent;
                }
                else if (cmd == NetworkEvent.Type.Data)
                {
                    // When the pong message is received we calculate the ping time and disconnect
                    m_lastPingTime = (int)((Time.fixedTime - m_pendingPing.time) * 1000);
                    m_clientToServerConnection.Disconnect(m_ClientDriver);
                    m_clientToServerConnection = default(NetworkConnection);
                    UnityEngine.Debug.Log($"Ping: {m_lastPingTime}");
                }
                else if (cmd == NetworkEvent.Type.Disconnect)
                {
                    // If the server disconnected us we clear out connection
                    m_clientToServerConnection = default(NetworkConnection);
                }
            }
        }

        void ServerUpdate()
        {
            // Update the NetworkDriver. It schedules a job so we must wait for that job with Complete
            m_ServerDriver.ScheduleUpdate().Complete();

            // Accept all new connections
            while (true)
            {
                var con = m_ServerDriver.Accept();
                // "Nothing more to accept" is signaled by returning an invalid connection from accept
                if (!con.IsCreated)
                    break;
                m_serverConnections.Add(con);
            }

            for (int i = 0; i < m_serverConnections.Length; ++i)
            {
                DataStreamReader strm;
                NetworkEvent.Type cmd;
                // Pop all events for the connection
                while ((cmd = m_ServerDriver.PopEventForConnection(m_serverConnections[i], out strm)) != NetworkEvent.Type.Empty)
                {
                    if (cmd == NetworkEvent.Type.Data)
                    {
                        // For ping requests we reply with a pong message
                        int id = strm.ReadInt();
                        // Create a temporary DataStreamWriter to keep our serialized pong message
                        if (m_ServerDriver.BeginSend(m_serverConnections[i], out var pongData) == 0)
                        {
                            pongData.WriteInt(id);
                            // Send the pong message with the same id as the ping
                            m_ServerDriver.EndSend(pongData);
                        }
                    }
                    else if (cmd == NetworkEvent.Type.Disconnect)
                    {
                        // This connection no longer exist, remove it from the list
                        // The next iteration will operate on the new connection we swapped in so as long as it exist the
                        // loop can continue
                        m_serverConnections.RemoveAtSwapBack(i);
                        if (i >= m_serverConnections.Length)
                            break;
                    }
                }
            }
        }
    }
}
