using Microsoft.AspNetCore.Routing.Constraints;
using Microsoft.AspNetCore.Server.HttpSys;
using System.Net.Sockets;
using System.Text;

namespace NotPogo
{
    public static class TCPProtocol
    {
        /// <summary>
        /// Reverses Pogo XOR Obfucscation And Reads Key Value Pairs From TCP Stream
        /// </summary>
        /// <param name="stream">The Stream Associated With The Received Message</param>
        /// <returns>A Dictionary Containing A Command And Data Pair (Key = Command, Value = Data)</returns>
        public static async Task<Dictionary<string, string>> ReadMessageAsync(NetworkStream stream)
        {
            int len = await ReadUInt16BEAsync(stream);
            byte[] payload = await ReadExactAsync(stream, len);

            for (int i = 0; i < payload.Length; i++)
                payload[i] = (byte)(payload[i] ^ 0xFF);

            var props = new Dictionary<string, string>();
            using var ms = new MemoryStream(payload);
            while(ms.Position < ms.Length)
            {
                string key = ReadJavaUtf(ms);
                string value = ReadJavaUtf(ms);
                props[key] = value;
            }
            return props;
        }

        /// <summary>
        /// Writes Key Value Pairs To A TCP Stream
        /// </summary>
        /// <param name="stream">The Stream You Wish To Write To</param>
        /// <param name="props">The Properties You Wish To Write To The Stream</param>
        /// <returns>The Completed Task</returns>
        public static async Task WriteMessageAsync(NetworkStream stream, IEnumerable<KeyValuePair<string, string>> props)
        {
            using var payload = new MemoryStream();
            foreach(var kv in props)
            {
                WriteJavaUtf(payload, kv.Key);
                WriteJavaUtf(payload, kv.Value);
            }

            byte[] bytes = payload.ToArray();
            for (int i = 0; i < bytes.Length; i++)
                bytes[i] = (byte)(bytes[i] ^ 0xFF);

            await WriteUint16BEAsync(stream, (ushort)bytes.Length);
            await stream.WriteAsync(bytes);
            await stream.FlushAsync();
        }

        private static async Task<int> ReadUInt16BEAsync(NetworkStream stream)
        {
            byte[] buf = await ReadExactAsync(stream, 2);
            return (buf[0] << 8) | buf[1];
        }

        private static async Task<byte[]> ReadExactAsync(NetworkStream stream, int count)
        {
            if (count < 0)
                throw new IOException($"Invalid Frame Length: {count}");

            byte[] buf = new byte[count];
            int off = 0;
            while (off < count)
            {
                int n = await stream.ReadAsync(buf.AsMemory(off, count - off));
                if (n <= 0)
                    throw new EndOfStreamException($"Socket Closed Mid-Message ({off}/{count} bytes read)");
                off += n;
            }
            return buf;
        }

        private static string ReadJavaUtf(Stream s)
        {
            int len = (s.ReadByte() << 8) | s.ReadByte();
            byte[] bytes = new byte[len];
            int off = 0;

            while (off < len)
            {
                int n = s.Read(bytes, off, len - off);
                if (n <= 0)
                    throw new EndOfStreamException();
                off += n;
            }

            var sb = new StringBuilder();
            int i = 0;
            while (i < bytes.Length)
            {
                int c = bytes[i] & 0xFF;
                if (c >> 4 <= 7)
                {
                    sb.Append((char)c);
                    i += 1;
                }
                else if (c >> 4 == 12 || c >> 4 == 13)
                {
                    sb.Append((char)(c & 0x1F << 6) | (bytes[i + 1] & 0x3F));
                    i += 2;
                }
                else
                {
                    sb.Append((char)((c & 0x0F) << 12) | ((bytes[i + 1] & 0x3F) << 6) | (bytes[i + 2] & 0x3F));
                    i += 3;
                }
            }
            return sb.ToString();
        }

        private static async Task WriteUint16BEAsync(NetworkStream stream, ushort length)
        {
            await stream.WriteAsync(new byte[] { (byte)(length >> 8), (byte)length });
        }

        private static void WriteJavaUtf(Stream s, string str)
        {
            using var buf = new MemoryStream();
            foreach(char ch in str)
            {
                if (ch >= 0x001 && ch <= 0x007F)
                    buf.WriteByte((byte)ch);
                else if (ch == 0 || ch > 0x07FF)
                {
                    buf.WriteByte((byte)(0xE0 | ((ch >> 12) & 0x0F)));
                    buf.WriteByte((byte)(0x80 | ((ch >> 6) & 0x3F)));
                    buf.WriteByte((byte)(0x80 | (ch & 0x3F)));
                }
                else
                {
                    buf.WriteByte((byte)(0x0C | ((ch >> 6) & 0x1F)));
                    buf.WriteByte((byte)(0x80 | (ch & 0x3F)));
                }
            }
            byte[] bytes = buf.ToArray();
            s.WriteByte((byte)(bytes.Length >> 8));
            s.WriteByte((byte)bytes.Length);
            s.Write(bytes, 0, bytes.Length);
        }
    }
}