using System;
using System.Collections.Generic;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Networking.Transport;
using UnityEngine;
using Logger = RebootKit.Engine.Foundation.Logger;

namespace RebootKit.Engine.Network {
    // @NOTE: This won't probably be suited for production use.
    public class UnityNetworkManager : INetworkManager {
        static readonly Logger s_Logger = new Logger(nameof(UnityNetworkManager));
        
        const ulong k_ServerClientID = 0;

        NetworkDriver m_Driver;

        // @MARK: Server specific stuff
        bool m_IsServer;
        readonly Dictionary<ulong, NetworkConnection> m_ServerConnections = new Dictionary<ulong, NetworkConnection>();
        ulong m_ClientIDCounter = 1;

        bool m_IsClient;
        NetworkConnection m_ClientConnection;

        public INetworkManagerDelegate Delegate { get; set; }

        public ulong LocalClientID { get; private set; }
        
        public NetworkManagerStats Stats { get; } = new NetworkManagerStats();

        public ushort Port = 7777;

        public UnityNetworkManager() {
            m_Driver = NetworkDriver.Create();
            if (!m_Driver.IsCreated) {
                s_Logger.Error("Failed to create network driver");
                return;
            }

            m_IsServer = false;
            m_IsClient = false;

            LocalClientID = 0; // This should be set to a unique ID for the local client
        }

        public void Dispose() {
            if (m_Driver.IsCreated) {
                m_Driver.Dispose();
            }

            m_ServerConnections.Clear();
            m_ClientConnection = default;
        }
        
        float m_StatsTimer = 0.0f;

        public void Tick() {
            m_Driver.ScheduleUpdate().Complete();

            if (IsServer()) {
                ServerTick();
            }

            if (IsClient()) {
                ClientTick();
            }
            
            if (IsServer() || IsClient()) {
                m_StatsTimer -= Time.deltaTime;
                if (m_StatsTimer <= 0.0f) {
                    m_StatsTimer = 1.0f;
                    
                    Stats.ReliableBytesSentPerSecond = Stats.ReliableBytesSent;
                    Stats.UnreliableBytesSentPerSecond = Stats.UnreliableBytesSent;
                    Stats.BytesReceivedPerSecond = Stats.BytesReceived;

                    Stats.ReliableBytesSent = 0;
                    Stats.UnreliableBytesSent = 0;
                    Stats.BytesReceived = 0;
                }
            }
        }

        void ServerTick() {
            using NativeList<ulong> clientIDsToRemove = new NativeList<ulong>(Allocator.Temp);
            foreach ((ulong clientID, NetworkConnection connection) in m_ServerConnections) {
                if (!connection.IsCreated) {
                    clientIDsToRemove.Add(clientID);
                }
            }

            foreach (ulong clientID in clientIDsToRemove) {
                m_ServerConnections.Remove(clientID);
                Delegate?.OnClientDisconnected(clientID);
                s_Logger.Info($"Client {clientID} disconnected");
            }

            NetworkConnection incomingConnection;
            while ((incomingConnection = m_Driver.Accept()) != default) {
                ulong newClientID = m_ClientIDCounter++;
                m_ServerConnections.Add(newClientID, incomingConnection);

                s_Logger.Info("Connection accepted: " + newClientID);
                Delegate?.OnClientConnected(newClientID);
            }

            // @NOTE: Handle incoming messages from clients
            clientIDsToRemove.Clear();
            foreach ((ulong clientID, NetworkConnection connection) in m_ServerConnections) {
                DataStreamReader stream;
                NetworkEvent.Type cmd;
                while ((cmd = m_Driver.PopEventForConnection(connection, out stream)) !=
                       NetworkEvent.Type.Empty) {
                    if (cmd == NetworkEvent.Type.Data) {
                        NativeArray<byte> data = new NativeArray<byte>(stream.Length, Allocator.Temp);
                        stream.ReadBytes(data);

                        Stats.BytesReceived += (ulong) stream.Length;
                        Delegate?.OnMessageReceived(clientID, data);
                    } else if (cmd == NetworkEvent.Type.Disconnect) {
                        clientIDsToRemove.Add(clientID);
                    }
                }
            }

            // @TODO: Code duplication
            foreach (ulong clientID in clientIDsToRemove) {
                m_ServerConnections.Remove(clientID);
                Delegate?.OnClientDisconnected(clientID);
                s_Logger.Info($"Client {clientID} disconnected");
            }
        }

        void ClientTick() {
            DataStreamReader stream;
            NetworkEvent.Type cmd;
            while ((cmd = m_ClientConnection.PopEvent(m_Driver, out stream)) != NetworkEvent.Type.Empty) {
                if (cmd == NetworkEvent.Type.Data) {
                    NativeArray<byte> data = new NativeArray<byte>(stream.Length, Allocator.Temp);
                    stream.ReadBytes(data);
                    Stats.BytesReceived += (ulong) stream.Length;
                    Delegate?.OnMessageReceived(k_ServerClientID, data);
                } else if (cmd == NetworkEvent.Type.Disconnect) {
                    m_IsClient = false;
                    m_ClientConnection = default;
                    Delegate?.OnClientDisconnected(LocalClientID);
                    s_Logger.Info("Disconnected from server");
                } else if (cmd == NetworkEvent.Type.Connect) {
                    s_Logger.Info("Client connected");
                    Delegate?.OnClientConnected(LocalClientID);
                }
            }
        }

        //
        // @MARK: INetworkManager
        //
        public bool IsServer() {
            return m_IsServer;
        }

        public bool StartHost() {
            if (IsServer()) {
                s_Logger.Error("Server already started");
                return false;
            }

            Stats.Reset();

            NetworkEndpoint endpoint = NetworkEndpoint.AnyIpv4.WithPort(7777);

            if (m_Driver.Bind(endpoint) != 0) {
                s_Logger.Error($"Failed to bind to port {Port}");
                return false;
            }

            m_Driver.Listen();

            m_ClientConnection = m_Driver.Connect(NetworkEndpoint.LoopbackIpv4.WithPort(Port));
            
            if (!m_ClientConnection.IsCreated) {
                s_Logger.Error($"Failed to create client connection on port {Port}");
                return false;
            }

            LocalClientID = k_ServerClientID;
            m_ServerConnections.Add(LocalClientID, m_ClientConnection);

            m_IsServer = true;
            m_IsClient = true;
            
            Delegate?.OnServerStarted();
            Delegate?.OnClientConnected(LocalClientID);
            return true;
        }

        public void StopHost() { throw new NotImplementedException(); }

        public bool IsClient() {
            return m_IsClient;
        }

        public bool StartClient() {
            if (IsClient()) {
                s_Logger.Error("Client already started");
                return false;
            }

            Stats.Reset();

            NetworkEndpoint endpoint = NetworkEndpoint.LoopbackIpv4.WithPort(Port);
            m_ClientConnection = m_Driver.Connect(endpoint);
            
            if (!m_ClientConnection.IsCreated) {
                s_Logger.Error($"Failed to connect to server on port {Port}");
                return false;
            }
            
            m_IsClient = true;
            Delegate?.OnClientStarted();
            return true;
        }

        public void Disconnect() { throw new NotImplementedException(); }

        public void Send(ulong clientID, NativeSlice<byte> data, SendMode mode) {
            unsafe {
                void* ptr = data.GetUnsafePtr();
                Send(clientID, (byte*)ptr, data.Length, mode);
            }
        }

        public unsafe void Send(ulong clientID, byte* data, int length, SendMode mode) {
            if (IsServer()) {
                if (clientID == LocalClientID) {
                    NativeArray<byte> messageData = new NativeArray<byte>(length, Allocator.Temp);
                    UnsafeUtility.MemCpy(messageData.GetUnsafePtr(), data, length);
                    Delegate?.OnMessageReceived(clientID, messageData);
                    return;
                }

                if (m_ServerConnections.TryGetValue(clientID, out NetworkConnection value)) {
                    if (!value.IsCreated) {
                        s_Logger.Error($"Client {clientID} connection is not created. Cannot send message.");
                        return;
                    }

                    m_Driver.BeginSend(value, out DataStreamWriter writer);
                    writer.WriteBytesUnsafe(data, length);
                    m_Driver.EndSend(writer);

                    if (mode == SendMode.Reliable) {
                        Stats.ReliableBytesSent += (ulong) length;
                    } else if (mode == SendMode.Unreliable){
                        Stats.UnreliableBytesSent += (ulong) length;
                    }
                } else {
                    s_Logger.Error($"Client {clientID} not found. Cannot send message.");
                }
            } else if (IsClient()) {
                if (clientID != k_ServerClientID) {
                    s_Logger.Error("Client ID mismatch. Cannot send message to another client from client context.");
                    return;
                }
                
                if (!m_ClientConnection.IsCreated) {
                    s_Logger.Error("Client connection is not created. Cannot send message.");
                    return;
                }

                if (m_Driver.BeginSend(m_ClientConnection, out DataStreamWriter writer, length) == 0) {
                    writer.WriteBytesUnsafe(data, length);

                    if (m_Driver.EndSend(writer) < 0) {
                        s_Logger.Error("Failed to send message to server.");
                    }

                    if (mode == SendMode.Reliable) {
                        Stats.ReliableBytesSent += (ulong) length;
                    } else if (mode == SendMode.Unreliable){
                        Stats.UnreliableBytesSent += (ulong) length;
                    }
                } else {
                    s_Logger.Error($"Failed to begin sending message to server with length {length}.");
                }
            }
        }
    }
}