using System;
using System.Text;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Assertions;
using UnityEngine.Pool;

namespace RebootKit.Engine.Network {
    // @NOTE: Data is written in a linear fashion, so the position is always at the end of the written data.
    //        We are writting everything in little-endian format.
    public struct NetworkBufferWriter : IDisposable {
        class WriterHandle {
            public NativeArray<byte> Data;
            public bool IsOwner; // Indicates if this handle owns the data and should dispose it.
            public int Position;
            public int Capacity;
        }

        static readonly IObjectPool<WriterHandle> s_WriterPool = new ObjectPool<WriterHandle>(
             () => new WriterHandle(),
             _ => { },
             handle => {
                 if (handle.Data.IsCreated && handle.IsOwner) {
                     handle.Data.Dispose();
                 }

                 handle.Data = default;
                 handle.Position = 0;
                 handle.Capacity = 0;
             },
             handle => {
                 if (handle.Data.IsCreated && handle.IsOwner) {
                     handle.Data.Dispose();
                 }
             },
             true,
             256
            );

        WriterHandle m_Handle;
        
        public int Position {
            get {
                return m_Handle.Position;
            }

            set {
                Assert.IsTrue(value >= 0 && value <= m_Handle.Capacity, "Position must be within the bounds of the buffer.");
                m_Handle.Position = value;
            }
        }

        public NetworkBufferWriter(int capacity, Allocator allocator) {
            m_Handle = s_WriterPool.Get();
            m_Handle.Data = new NativeArray<byte>(capacity, allocator);
            m_Handle.IsOwner = true;
            m_Handle.Capacity = capacity;
            m_Handle.Position = 0;
        }

        public NetworkBufferWriter(NativeArray<byte> buffer, int position) {
            m_Handle = s_WriterPool.Get();
            m_Handle.Data = buffer;
            m_Handle.IsOwner = false;
            m_Handle.Capacity = buffer.Length;
            m_Handle.Position = position;
        }

        public void Dispose() {
            if (m_Handle != null) {
                s_WriterPool.Release(m_Handle);
                m_Handle = null;
            }
        }

        public bool WillFit(int size) {
            return m_Handle.Position + size <= m_Handle.Capacity;
        }

        public void Write(byte value) {
            if (m_Handle.Position >= m_Handle.Capacity) {
                throw new InvalidOperationException("Buffer overflow: Cannot write beyond capacity.");
            }
            m_Handle.Data[m_Handle.Position++] = value;
        }

        public void Write(byte[] values) {
            Assert.IsNotNull(values, "Trying to write null byte array to the buffer.");
            Assert.IsTrue(WillFit(values.Length), "Buffer overflow: Cannot write beyond capacity.");

            if (values.Length == 0) {
                return;
            }

            for (int i = 0; i < values.Length; i++) {
                m_Handle.Data[m_Handle.Position++] = values[i];
            }
        }

        public void Write(NativeArray<byte> values) {
            Assert.IsTrue(values.IsCreated, "Trying to write uncreated NativeArray to the buffer.");
            Assert.IsTrue(WillFit(values.Length), "Buffer overflow: Cannot write beyond capacity.");

            if (values.Length == 0) {
                return;
            }

            for (int i = 0; i < values.Length; i++) {
                m_Handle.Data[m_Handle.Position++] = values[i];
            }
        }

        public void Write(int value) {
            Assert.IsTrue(sizeof(int) == 4, "Size of int must be 4 bytes.");
            Assert.IsTrue(WillFit(sizeof(int)), "Buffer overflow: Cannot write beyond capacity.");

            if (BitConverter.IsLittleEndian) {
                Write((byte) (value & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) ((value >> 16) & 0xFF));
                Write((byte) ((value >> 24) & 0xFF));
            } else {
                Write((byte) ((value >> 24) & 0xFF));
                Write((byte) ((value >> 16) & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) (value & 0xFF));
            }
        }

        public void Write(short value) {
            Assert.IsTrue(sizeof(short) == 2, "Size of short must be 2 bytes.");
            Assert.IsTrue(WillFit(sizeof(short)), "Buffer overflow: Cannot write beyond capacity.");

            if (BitConverter.IsLittleEndian) {
                Write((byte) (value & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
            } else {
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) (value & 0xFF));
            }
        }

        public void Write(ushort value) {
            Assert.IsTrue(sizeof(ushort) == 2, "Size of ushort must be 2 bytes.");
            Assert.IsTrue(WillFit(sizeof(ushort)), "Buffer overflow: Cannot write beyond capacity.");

            if (BitConverter.IsLittleEndian) {
                Write((byte) (value & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
            } else {
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) (value & 0xFF));
            }
        }

        public void Write(long value) {
            Assert.IsTrue(sizeof(long) == 8, "Size of long must be 8 bytes.");
            Assert.IsTrue(WillFit(sizeof(long)), "Buffer overflow: Cannot write beyond capacity.");

            if (BitConverter.IsLittleEndian) {
                Write((byte) (value & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) ((value >> 16) & 0xFF));
                Write((byte) ((value >> 24) & 0xFF));
                Write((byte) ((value >> 32) & 0xFF));
                Write((byte) ((value >> 40) & 0xFF));
                Write((byte) ((value >> 48) & 0xFF));
                Write((byte) ((value >> 56) & 0xFF));
            } else {
                Write((byte) ((value >> 56) & 0xFF));
                Write((byte) ((value >> 48) & 0xFF));
                Write((byte) ((value >> 40) & 0xFF));
                Write((byte) ((value >> 32) & 0xFF));
                Write((byte) ((value >> 24) & 0xFF));
                Write((byte) ((value >> 16) & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) (value & 0xFF));
            }
        }

        public void Write(ulong value) {
            Assert.IsTrue(sizeof(ulong) == 8, "Size of ulong must be 8 bytes.");
            Assert.IsTrue(WillFit(sizeof(ulong)), $"Buffer overflow: Cannot write beyond capacity. Current position: {m_Handle.Position}, Capacity: {m_Handle.Capacity}");

            if (BitConverter.IsLittleEndian) {
                Write((byte) (value & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) ((value >> 16) & 0xFF));
                Write((byte) ((value >> 24) & 0xFF));
                Write((byte) ((value >> 32) & 0xFF));
                Write((byte) ((value >> 40) & 0xFF));
                Write((byte) ((value >> 48) & 0xFF));
                Write((byte) ((value >> 56) & 0xFF));
            } else {
                Write((byte) ((value >> 56) & 0xFF));
                Write((byte) ((value >> 48) & 0xFF));
                Write((byte) ((value >> 40) & 0xFF));
                Write((byte) ((value >> 32) & 0xFF));
                Write((byte) ((value >> 24) & 0xFF));
                Write((byte) ((value >> 16) & 0xFF));
                Write((byte) ((value >> 8) & 0xFF));
                Write((byte) (value & 0xFF));
            }
        }

        public void Write(float value) {
            Assert.IsTrue(sizeof(float) == 4, "Size of float must be 4 bytes.");
            Assert.IsTrue(WillFit(sizeof(float)), "Buffer overflow: Cannot write beyond capacity.");

            unsafe {
                byte* bytes = (byte*) &value;
                Write(bytes[0]);
                Write(bytes[1]);
                Write(bytes[2]);
                Write(bytes[3]);
            }
        }

        public void Write(bool value) {
            Assert.IsTrue(WillFit(1), "Buffer overflow: Cannot write beyond capacity.");
            Write((byte) (value ? 1 : 0));
        }

        public void Write(Vector2 value) {
            Assert.IsTrue(WillFit(sizeof(float) * 2), "Buffer overflow: Cannot write beyond capacity.");

            Write(value.x);
            Write(value.y);
        }

        public void Write(Vector3 value) {
            Assert.IsTrue(WillFit(sizeof(float) * 3), "Buffer overflow: Cannot write beyond capacity.");

            Write(value.x);
            Write(value.y);
            Write(value.z);
        }

        public void Write(Vector4 value) {
            Assert.IsTrue(WillFit(sizeof(float) * 4), "Buffer overflow: Cannot write beyond capacity.");

            Write(value.x);
            Write(value.y);
            Write(value.z);
            Write(value.w);
        }

        public void Write(Quaternion value) {
            Assert.IsTrue(WillFit(sizeof(float) * 4), "Buffer overflow: Cannot write beyond capacity.");

            Write(value.x);
            Write(value.y);
            Write(value.z);
            Write(value.w);
        }

        public void Write(FixedString32Bytes value) {
            Assert.IsTrue(WillFit(32));

            for (int i = 0; i < 32; i++) {
                if (i < value.Length) {
                    Write(value[i]);
                } else {
                    Write((byte) 0); // Fill with zero if the string is shorter than 32 bytes
                }
            }
        }
        
        public void Write(FixedString64Bytes value) {
            Assert.IsTrue(WillFit(64));

            for (int i = 0; i < 64; i++) {
                if (i < value.Length) {
                    Write(value[i]);
                } else {
                    Write((byte) 0); // Fill with zero if the string is shorter than 64 bytes
                }
            }
        }
        
        public void Write(FixedString128Bytes value) {
            Assert.IsTrue(WillFit(128));

            for (int i = 0; i < 128; i++) {
                if (i < value.Length) {
                    Write(value[i]);
                } else {
                    Write((byte) 0); // Fill with zero if the string is shorter than 128 bytes
                }
            }
        }
        
        public void Write(FixedString512Bytes value) {
            Assert.IsTrue(WillFit(512));

            for (int i = 0; i < 512; i++) {
                if (i < value.Length) {
                    Write(value[i]);
                } else {
                    Write((byte) 0); // Fill with zero if the string is shorter than 512 bytes
                }
            }
        }

        public void Write(FixedString4096Bytes value) {
            Assert.IsTrue(WillFit(4096));

            for (int i = 0; i < 4096; i++) {
                if (i < value.Length) {
                    Write(value[i]);
                } else {
                    Write((byte) 0); // Fill with zero if the string is shorter than 4096 bytes
                }
            }
        }

        // @NOTE: Writes ascii characters as bytes, and prepends the length as a ushort.
        public static int GetStringWriteLength(string value) {
            if (value.Length > ushort.MaxValue) {
                return sizeof(ushort);
            }
            
            if (string.IsNullOrEmpty(value)) {
                return sizeof(ushort);
            }
            
            // Length prefix
            int length = sizeof(ushort); 
            
            for (int i = 0; i < value.Length; i++) {
                length += sizeof(char);
            }

            return length;
        }

        public void Write(string value) {
            if (value.Length > ushort.MaxValue) {
                Write((ushort) 0);
                return;
            }
            
            if (string.IsNullOrEmpty(value)) {
                Write((ushort) 0);
                return;
            }
            
            Write((ushort)value.Length);
            
            for (int i = 0; i < value.Length; i++) {
                Write((byte)value[i]);
            }
        }
    }
}