using System.Reflection;
using System.Security.Cryptography;
using System.Text;
using OF_DL.Crypto;
using ProtoBuf;

namespace OF_DL.Widevine;

public class CDM
{
    private static Dictionary<string, CDMDevice> Devices { get; } = new()
    {
        [Constants.DEVICE_NAME] = new CDMDevice(Constants.DEVICE_NAME)
    };

    private static Dictionary<string, Session> Sessions { get; } = new();

    private static byte[] CheckPSSH(string psshB64)
    {
        byte[] systemID = new byte[] { 237, 239, 139, 169, 121, 214, 74, 206, 163, 200, 39, 220, 213, 29, 33, 237 };

        if (psshB64.Length % 4 != 0)
        {
            psshB64 = psshB64.PadRight(psshB64.Length + (4 - psshB64.Length % 4), '=');
        }

        byte[] pssh = Convert.FromBase64String(psshB64);

        if (pssh.Length < 30)
        {
            return pssh;
        }

        if (!pssh[12..28].SequenceEqual(systemID))
        {
            List<byte> newPssh = new() { 0, 0, 0 };
            newPssh.Add((byte)(32 + pssh.Length));
            newPssh.AddRange(Encoding.UTF8.GetBytes("pssh"));
            newPssh.AddRange(new byte[] { 0, 0, 0, 0 });
            newPssh.AddRange(systemID);
            newPssh.AddRange(new byte[] { 0, 0, 0, 0 });
            newPssh[31] = (byte)pssh.Length;
            newPssh.AddRange(pssh);

            return newPssh.ToArray();
        }

        return pssh;
    }

    public static string? OpenSession(string initDataB64, string deviceName, bool offline = false, bool raw = false)
    {
        byte[] initData = CheckPSSH(initDataB64);

        CDMDevice device = Devices[deviceName];

        byte[] sessionId = new byte[16];

        if (device.IsAndroid)
        {
            string randHex = "";

            Random rand = new();
            string choice = "ABCDEF0123456789";
            for (int i = 0; i < 16; i++)
            {
                randHex += choice[rand.Next(16)];
            }

            string counter = "01";
            string rest = "00000000000000";
            sessionId = Encoding.ASCII.GetBytes(randHex + counter + rest);
        }
        else
        {
            Random rand = new();
            rand.NextBytes(sessionId);
        }

        Session session;
        dynamic? parsedInitData = ParseInitData(initData);

        if (parsedInitData != null)
        {
            session = new Session(sessionId, parsedInitData, device, offline);
        }
        else if (raw)
        {
            session = new Session(sessionId, initData, device, offline);
        }
        else
        {
            return null;
        }

        Sessions.Add(BytesToHex(sessionId), session);

        return BytesToHex(sessionId);
    }

    private static WidevineCencHeader? ParseInitData(byte[] initData)
    {
        WidevineCencHeader cencHeader;

        try
        {
            cencHeader = Serializer.Deserialize<WidevineCencHeader>(new MemoryStream(initData[32..]));
        }
        catch
        {
            try
            {
                //needed for HBO Max

                PsshBox psshBox = PsshBox.FromByteArray(initData);
                cencHeader = Serializer.Deserialize<WidevineCencHeader>(new MemoryStream(psshBox.Data));
            }
            catch
            {
                //Logger.Verbose("Unable to parse, unsupported init data format");
                return null;
            }
        }

        return cencHeader;
    }

    public static bool CloseSession(string sessionId)
    {
        //Logger.Debug($"CloseSession(session_id={BytesToHex(sessionId)})");
        //Logger.Verbose("Closing CDM session");

        if (Sessions.ContainsKey(sessionId))
        {
            Sessions.Remove(sessionId);
            //Logger.Verbose("CDM session closed");
            return true;
        }

        //Logger.Info($"Session {sessionId} not found");
        return false;
    }

    public static bool SetServiceCertificate(string sessionId, byte[] certData)
    {
        //Logger.Debug($"SetServiceCertificate(sessionId={BytesToHex(sessionId)}, cert={certB64})");
        //Logger.Verbose($"Setting service certificate");

        if (!Sessions.ContainsKey(sessionId))
        {
            //Logger.Error("Session ID doesn't exist");
            return false;
        }

        SignedMessage signedMessage = new();

        try
        {
            signedMessage = Serializer.Deserialize<SignedMessage>(new MemoryStream(certData));
        }
        catch
        {
            //Logger.Warn("Failed to parse cert as SignedMessage");
        }

        SignedDeviceCertificate serviceCertificate;
        try
        {
            try
            {
                //Logger.Debug("Service cert provided as signedmessage");
                serviceCertificate =
                    Serializer.Deserialize<SignedDeviceCertificate>(new MemoryStream(signedMessage.Msg));
            }
            catch
            {
                //Logger.Debug("Service cert provided as signeddevicecertificate");
                serviceCertificate = Serializer.Deserialize<SignedDeviceCertificate>(new MemoryStream(certData));
            }
        }
        catch
        {
            //Logger.Error("Failed to parse service certificate");
            return false;
        }

        Sessions[sessionId].ServiceCertificate = serviceCertificate;
        Sessions[sessionId].PrivacyMode = true;

        return true;
    }

    public static byte[]? GetLicenseRequest(string sessionId)
    {
        //Logger.Debug($"GetLicenseRequest(sessionId={BytesToHex(sessionId)})");
        //Logger.Verbose($"Getting license request");

        if (!Sessions.ContainsKey(sessionId))
        {
            //Logger.Error("Session ID doesn't exist");
            return null;
        }

        Session session = Sessions[sessionId];

        //Logger.Debug("Building license request");

        dynamic licenseRequest;

        if (session.InitData is WidevineCencHeader)
        {
            licenseRequest = new SignedLicenseRequest
            {
                Type = SignedLicenseRequest.MessageType.LicenseRequest,
                Msg = new LicenseRequest
                {
                    Type = LicenseRequest.RequestType.New,
                    KeyControlNonce = 1093602366,
                    ProtocolVersion = ProtocolVersion.Current,
                    ContentId = new LicenseRequest.ContentIdentification
                    {
                        CencId = new LicenseRequest.ContentIdentification.Cenc
                        {
                            LicenseType = session.Offline ? LicenseType.Offline : LicenseType.Default,
                            RequestId = session.SessionId,
                            Pssh = session.InitData
                        }
                    }
                }
            };
        }
        else
        {
            licenseRequest = new SignedLicenseRequestRaw
            {
                Type = SignedLicenseRequestRaw.MessageType.LicenseRequest,
                Msg = new LicenseRequestRaw
                {
                    Type = LicenseRequestRaw.RequestType.New,
                    KeyControlNonce = 1093602366,
                    ProtocolVersion = ProtocolVersion.Current,
                    ContentId = new LicenseRequestRaw.ContentIdentification
                    {
                        CencId = new LicenseRequestRaw.ContentIdentification.Cenc
                        {
                            LicenseType = session.Offline ? LicenseType.Offline : LicenseType.Default,
                            RequestId = session.SessionId,
                            Pssh = session.InitData
                        }
                    }
                }
            };
        }

        if (session.PrivacyMode)
        {
            //Logger.Debug("Privacy mode & serivce certificate loaded, encrypting client id");

            EncryptedClientIdentification encryptedClientIdProto = new();

            //Logger.Debug("Unencrypted client id " + Utils.SerializeToString(clientId));

            using MemoryStream memoryStream = new();
            Serializer.Serialize(memoryStream, session.Device.ClientID);
            byte[] data = Padding.AddPKCS7Padding(memoryStream.ToArray(), 16);

            using Aes aes = Aes.Create();
            aes.BlockSize = 128;
            aes.Padding = PaddingMode.PKCS7;
            aes.Mode = CipherMode.CBC;

            aes.GenerateKey();
            aes.GenerateIV();

            using MemoryStream mstream = new();
            using CryptoStream cryptoStream = new(mstream, aes.CreateEncryptor(aes.Key, aes.IV),
                CryptoStreamMode.Write);
            cryptoStream.Write(data, 0, data.Length);
            encryptedClientIdProto.EncryptedClientId = mstream.ToArray();

            using RSACryptoServiceProvider RSA = new();
            RSA.ImportRSAPublicKey(session.ServiceCertificate.DeviceCertificate.PublicKey, out int _);
            encryptedClientIdProto.EncryptedPrivacyKey = RSA.Encrypt(aes.Key, RSAEncryptionPadding.OaepSHA1);
            encryptedClientIdProto.EncryptedClientIdIv = aes.IV;
            encryptedClientIdProto.ServiceId =
                Encoding.UTF8.GetString(session.ServiceCertificate.DeviceCertificate.ServiceId);
            encryptedClientIdProto.ServiceCertificateSerialNumber =
                session.ServiceCertificate.DeviceCertificate.SerialNumber;

            licenseRequest.Msg.EncryptedClientId = encryptedClientIdProto;
        }
        else
        {
            licenseRequest.Msg.ClientId = session.Device.ClientID;
        }

        //Logger.Debug("Signing license request");

        using (MemoryStream memoryStream = new())
        {
            Serializer.Serialize(memoryStream, licenseRequest.Msg);
            byte[] data = memoryStream.ToArray();
            session.LicenseRequest = data;

            licenseRequest.Signature = session.Device.Sign(data);
        }

        //Logger.Verbose("License request created");

        byte[] requestBytes;
        using (MemoryStream memoryStream = new())
        {
            Serializer.Serialize(memoryStream, licenseRequest);
            requestBytes = memoryStream.ToArray();
        }

        Sessions[sessionId] = session;

        //Logger.Debug($"license request b64: {Convert.ToBase64String(requestBytes)}");
        return requestBytes;
    }

    public static void ProvideLicense(string sessionId, byte[] license)
    {
        //Logger.Debug($"ProvideLicense(sessionId={BytesToHex(sessionId)}, licenseB64={licenseB64})");
        //Logger.Verbose("Decrypting provided license");

        if (!Sessions.ContainsKey(sessionId))
        {
            throw new Exception("Session ID doesn't exist");
        }

        Session session = Sessions[sessionId];

        if (session.LicenseRequest == null)
        {
            throw new Exception("Generate a license request first");
        }

        SignedLicense signedLicense;
        try
        {
            signedLicense = Serializer.Deserialize<SignedLicense>(new MemoryStream(license));
        }
        catch
        {
            throw new Exception("Unable to parse license");
        }

        //Logger.Debug("License: " + Utils.SerializeToString(signedLicense));

        session.License = signedLicense;

        //Logger.Debug($"Deriving keys from session key");

        try
        {
            byte[] sessionKey = session.Device.Decrypt(session.License.SessionKey);

            if (sessionKey.Length != 16)
            {
                throw new Exception("Unable to decrypt session key");
            }

            session.SessionKey = sessionKey;
        }
        catch
        {
            throw new Exception("Unable to decrypt session key");
        }

        //Logger.Debug("Session key: " + BytesToHex(session.SessionKey));

        session.DerivedKeys = DeriveKeys(session.LicenseRequest, session.SessionKey);

        //Logger.Debug("Verifying license signature");

        byte[] licenseBytes;
        using (MemoryStream memoryStream = new())
        {
            Serializer.Serialize(memoryStream, signedLicense.Msg);
            licenseBytes = memoryStream.ToArray();
        }

        byte[] hmacHash = CryptoUtils.GetHMACSHA256Digest(licenseBytes, session.DerivedKeys.Auth1);

        if (!hmacHash.SequenceEqual(signedLicense.Signature))
        {
            throw new Exception("License signature mismatch");
        }

        foreach (License.KeyContainer key in signedLicense.Msg.Keys)
        {
            string type = key.Type.ToString();

            if (type == "Signing")
            {
                continue;
            }

            byte[] encryptedKey = key.Key;
            byte[] iv = key.Iv;
            byte[] keyId = key.Id;
            if (keyId == null)
            {
                keyId = Encoding.ASCII.GetBytes(key.Type.ToString());
            }

            using MemoryStream mstream = new();

            using Aes aes = Aes.Create();
            aes.Padding = PaddingMode.PKCS7;
            aes.Mode = CipherMode.CBC;

            using CryptoStream cryptoStream = new(mstream, aes.CreateDecryptor(session.DerivedKeys.Enc, iv),
                CryptoStreamMode.Write);
            cryptoStream.Write(encryptedKey, 0, encryptedKey.Length);
            byte[] decryptedKey = mstream.ToArray();

            List<string> permissions = [];
            if (type == "OperatorSession")
            {
                foreach (PropertyInfo perm in key._OperatorSessionKeyPermissions.GetType().GetProperties())
                {
                    if ((uint?)perm.GetValue(key._OperatorSessionKeyPermissions) == 1)
                    {
                        permissions.Add(perm.Name);
                    }
                }
            }

            session.ContentKeys.Add(new ContentKey
            {
                KeyID = keyId, Type = type, Bytes = decryptedKey, Permissions = permissions
            });
        }

        //Logger.Debug($"Key count: {session.Keys.Count}");

        Sessions[sessionId] = session;

        //Logger.Verbose("Decrypted all keys");
    }

    public static DerivedKeys DeriveKeys(byte[] message, byte[] key)
    {
        byte[] encKeyBase = Encoding.UTF8.GetBytes("ENCRYPTION").Concat(new byte[] { 0x0 }).Concat(message)
            .Concat(new byte[] { 0x0, 0x0, 0x0, 0x80 }).ToArray();
        byte[] authKeyBase = Encoding.UTF8.GetBytes("AUTHENTICATION").Concat(new byte[] { 0x0 }).Concat(message)
            .Concat(new byte[] { 0x0, 0x0, 0x2, 0x0 }).ToArray();

        byte[] encKey = new byte[] { 0x01 }.Concat(encKeyBase).ToArray();
        byte[] authKey1 = new byte[] { 0x01 }.Concat(authKeyBase).ToArray();
        byte[] authKey2 = new byte[] { 0x02 }.Concat(authKeyBase).ToArray();
        byte[] authKey3 = new byte[] { 0x03 }.Concat(authKeyBase).ToArray();
        byte[] authKey4 = new byte[] { 0x04 }.Concat(authKeyBase).ToArray();

        byte[] encCmacKey = CryptoUtils.GetCMACDigest(encKey, key);
        byte[] authCmacKey1 = CryptoUtils.GetCMACDigest(authKey1, key);
        byte[] authCmacKey2 = CryptoUtils.GetCMACDigest(authKey2, key);
        byte[] authCmacKey3 = CryptoUtils.GetCMACDigest(authKey3, key);
        byte[] authCmacKey4 = CryptoUtils.GetCMACDigest(authKey4, key);

        byte[] authCmacCombined1 = authCmacKey1.Concat(authCmacKey2).ToArray();
        byte[] authCmacCombined2 = authCmacKey3.Concat(authCmacKey4).ToArray();

        return new DerivedKeys { Auth1 = authCmacCombined1, Auth2 = authCmacCombined2, Enc = encCmacKey };
    }

    public static List<ContentKey> GetKeys(string sessionId)
    {
        if (Sessions.ContainsKey(sessionId))
        {
            return Sessions[sessionId].ContentKeys;
        }

        throw new Exception("Session not found");
    }

    private static string BytesToHex(byte[] data) => BitConverter.ToString(data).Replace("-", "");
}

/*
    public static List<string> ProvideLicense(string requestB64, string licenseB64)
    {
        byte[] licenseRequest;

        var request = Serializer.Deserialize<SignedLicenseRequest>(new MemoryStream(Convert.FromBase64String(requestB64)));

        using (var ms = new MemoryStream())
        {
            Serializer.Serialize(ms, request.Msg);
            licenseRequest = ms.ToArray();
        }

        SignedLicense signedLicense;
        try
        {
            signedLicense = Serializer.Deserialize<SignedLicense>(new MemoryStream(Convert.FromBase64String(licenseB64)));
        }
        catch
        {
            return null;
        }

        byte[] sessionKey;
        try
        {

            sessionKey = Controllers.Adapter.OaepDecrypt(Convert.ToBase64String(signedLicense.SessionKey));

            if (sessionKey.Length != 16)
            {
                return null;
            }
        }
        catch
        {
            return null;
        }

        byte[] encKeyBase = Encoding.UTF8.GetBytes("ENCRYPTION").Concat(new byte[] { 0x0, }).Concat(licenseRequest).Concat(new byte[] { 0x0, 0x0, 0x0, 0x80 }).ToArray();

        byte[] encKey = new byte[] { 0x01 }.Concat(encKeyBase).ToArray();

        byte[] encCmacKey = GetCmacDigest(encKey, sessionKey);

        byte[] encryptionKey = encCmacKey;

        List<string> keys = new List<string>();

        foreach (License.KeyContainer key in signedLicense.Msg.Keys)
        {
            string type = key.Type.ToString();
            if (type == "Signing")
            {
                continue;
            }

            byte[] keyId;
            byte[] encryptedKey = key.Key;
            byte[] iv = key.Iv;
            keyId = key.Id;
            if (keyId == null)
            {
                keyId = Encoding.ASCII.GetBytes(key.Type.ToString());
            }

            byte[] decryptedKey;

            using MemoryStream mstream = new MemoryStream();
            using AesCryptoServiceProvider aesProvider = new AesCryptoServiceProvider
            {
                Mode = CipherMode.CBC,
                Padding = PaddingMode.PKCS7
            };
            using CryptoStream cryptoStream = new CryptoStream(mstream, aesProvider.CreateDecryptor(encryptionKey, iv), CryptoStreamMode.Write);
            cryptoStream.Write(encryptedKey, 0, encryptedKey.Length);
            decryptedKey = mstream.ToArray();

            List<string> permissions = new List<string>();
            if (type == "OPERATOR_SESSION")
            {
                foreach (FieldInfo perm in key._OperatorSessionKeyPermissions.GetType().GetFields())
                {
                    if ((uint)perm.GetValue(key._OperatorSessionKeyPermissions) == 1)
                    {
                        permissions.Add(perm.Name);
                    }
                }
            }
            keys.Add(BitConverter.ToString(keyId).Replace("-","").ToLower() + ":" + BitConverter.ToString(decryptedKey).Replace("-", "").ToLower());
        }

        return keys;
    }*/
