using System;
using System.Security.Cryptography;
using static System.Buffers.Binary.BinaryPrimitives;

namespace PKHeX.Core;

/// <summary>
/// MemeCrypto V1 - The Original Series
/// </summary>
/// <remarks>
/// A variant of <see cref="SaveFile"/> encryption and obfuscation used in <see cref="GameVersion.Gen7"/>.
/// <br> The save file stores a dedicated block to contain a hash of the savedata, computed when the block is zeroed. </br>
/// <br> This signing logic is reused for other authentication; refer to <see cref="MemeKeyIndex"/>. </br>
/// <br> The save file first computes a SHA256 Hash over the block checksum region.
/// The logic then applies a SHA1 hash over the SHA256 hash result, encrypts it with a <see cref="MemeKey"/>, and signs it with an RSA private key in a non-straightforward manner. </br>
/// </remarks>
public static class MemeCrypto
{
    private const uint POKE = 0x454B4F50;

    public static bool VerifyMemePOKE(ReadOnlySpan<byte> input, out byte[] output)
    {
        if (input.Length < 0x60)
            throw new ArgumentException("Invalid POKE buffer!");
        var memeLen = input.Length - 8;
        var memeIndex = MemeKeyIndex.PokedexAndSaveFile;
        for (var i = input.Length - 8; i >= 0; i--)
        {
            if (ReadUInt32LittleEndian(input[i..]) != POKE)
                continue;

            var keyIndex = ReadInt32LittleEndian(input[(i+4)..]);
            if (!MemeKey.IsValidPokeKeyIndex(keyIndex))
                continue;

            memeLen = i;
            memeIndex = (MemeKeyIndex)keyIndex;
            break;
        }

        foreach (var len in new[] { memeLen, memeLen - 2 }) // Account for Pokédex QR Edge case
        {
            if (VerifyMemeData(input, out output, 0, len, memeIndex))
                return true;

            if (VerifyMemeData(input, out output, 0, len, MemeKeyIndex.PokedexAndSaveFile))
                return true;
        }

        output = Array.Empty<byte>();
        return false;
    }

    public static bool VerifyMemeData(ReadOnlySpan<byte> input, out byte[] output)
    {
        foreach (MemeKeyIndex keyIndex in Enum.GetValues(typeof(MemeKeyIndex)))
        {
            if (VerifyMemeData(input, out output, keyIndex))
                return true;
        }
        output = Array.Empty<byte>();
        return false;
    }

    public static bool VerifyMemeData(ReadOnlySpan<byte> input, out byte[] output, MemeKeyIndex keyIndex)
    {
        if (input.Length < 0x60)
        {
            output = Array.Empty<byte>();
            return false;
        }
        var key = new MemeKey(keyIndex);
        output = input.ToArray();

        var sigBuffer = key.RsaPublic(input[^0x60..]);
        using var sha1 = SHA1.Create();
        if (DecryptCompare(output, sigBuffer, key, sha1))
            return true;
        sigBuffer[0x0] |= 0x80;
        if (DecryptCompare(output, sigBuffer, key, sha1))
            return true;

        output = Array.Empty<byte>();
        return false;
    }

    private static bool DecryptCompare(byte[] output, ReadOnlySpan<byte> sigBuffer, MemeKey key, SHA1 sha1)
    {
        sigBuffer.CopyTo(output.AsSpan(output.Length - 0x60));
        key.AesDecrypt(output).CopyTo(output);
        // Check for 8-byte equality.
        var hash = sha1.ComputeHash(output, 0, output.Length - 0x8);
        var computed = ReadUInt64LittleEndian(hash.AsSpan());
        var existing = ReadUInt64LittleEndian(output.AsSpan(output.Length - 0x8));
        return computed == existing;
    }

    public static bool VerifyMemeData(ReadOnlySpan<byte> input, out byte[] output, int offset, int length)
    {
        var data = input.Slice(offset, length).ToArray();
        if (VerifyMemeData(data, out output))
        {
            var newOutput = input.ToArray();
            output.CopyTo(newOutput, offset);
            output = newOutput;
            return true;
        }
        output = Array.Empty<byte>();
        return false;
    }

    public static bool VerifyMemeData(ReadOnlySpan<byte> input, out byte[] output, int offset, int length, MemeKeyIndex keyIndex)
    {
        var data = input.Slice(offset, length);
        if (VerifyMemeData(data, out output, keyIndex))
        {
            var newOutput = input.ToArray();
            output.CopyTo(newOutput, offset);
            output = newOutput;
            return true;
        }
        output = Array.Empty<byte>();
        return false;
    }

    public static byte[] SignMemeData(ReadOnlySpan<byte> input, MemeKeyIndex keyIndex = MemeKeyIndex.PokedexAndSaveFile)
    {
        // Validate Input
        if (input.Length < 0x60)
            throw new ArgumentException("Cannot memesign a buffer less than 0x60 bytes in size!");
        var key = new MemeKey(keyIndex);
        if (!key.CanResign)
            throw new ArgumentException("Cannot sign with the specified memekey!");

        var output = input.ToArray();

        // Copy in the SHA1 signature
        using (var sha1 = SHA1.Create())
        {
            var hash = sha1.ComputeHash(output, 0, output.Length - 8);
            hash.AsSpan(0, 8).CopyTo(output.AsSpan(output.Length - 8, 8));
        }

        // Perform AES operations
        output = key.AesEncrypt(output);
        var sigBuffer = output.AsSpan(output.Length - 0x60, 0x60);
        sigBuffer[0] &= 0x7F;
        var signed = key.RsaPrivate(sigBuffer);
        signed.CopyTo(sigBuffer);
        return output;
    }

    /// <summary>
    /// Resigns save data.
    /// </summary>
    /// <param name="sav7">Save file data to resign</param>
    /// <returns>The resigned save data. Invalid input returns null.</returns>
    public static byte[] Resign7(ReadOnlySpan<byte> sav7)
    {
        if (sav7.Length is not (SaveUtil.SIZE_G7SM or SaveUtil.SIZE_G7USUM))
            throw new ArgumentException("Should not be using this for unsupported saves.");

        // Save Chunks are 0x200 bytes each; Memecrypto signature is 0x100 bytes into the 2nd to last chunk.
        var isUSUM = sav7.Length == SaveUtil.SIZE_G7USUM;
        var ChecksumTableOffset = sav7.Length - 0x200;
        var MemeCryptoOffset = isUSUM ? 0x6C100 : 0x6BB00;
        var ChecksumSignatureLength = isUSUM ? 0x150 : 0x140;
        const int MemeCryptoSignatureLength = 0x80;

        var result = sav7.ToArray();

        // Store current signature
        var oldSig = sav7.Slice(MemeCryptoOffset, MemeCryptoSignatureLength).ToArray();

        using var sha256 = SHA256.Create();
        var newSig = sha256.ComputeHash(result, ChecksumTableOffset, ChecksumSignatureLength);
        Span<byte> sigSpan = stackalloc byte[MemeCryptoSignatureLength];
        newSig.CopyTo(sigSpan);

        if (VerifyMemeData(oldSig, out var memeSig, MemeKeyIndex.PokedexAndSaveFile))
            memeSig.AsSpan()[0x20..0x80].CopyTo(sigSpan[0x20..]);

        SignMemeData(sigSpan).CopyTo(result, MemeCryptoOffset);
        return result;
    }
}