using System;
using System.Security.Cryptography;
using Org.BouncyCastle.Asn1;
using Org.BouncyCastle.Math;

namespace StellaOps.Cryptography;

public static class GostSignatureEncoding
{
    public static bool IsDer(ReadOnlySpan<byte> signature)
    {
        if (signature.Length < 2 || signature[0] != 0x30)
        {
            return false;
        }

        var lengthByte = signature[1];
        if ((lengthByte & 0x80) == 0)
        {
            var total = lengthByte + 2;
            return total == signature.Length;
        }

        var lengthBytes = lengthByte & 0x7F;
        if (lengthBytes is 0 or > 4 || signature.Length < 2 + lengthBytes)
        {
            return false;
        }

        var totalLength = 0;
        for (var i = 0; i < lengthBytes; i++)
        {
            totalLength = (totalLength << 8) | signature[2 + i];
        }

        return totalLength + 2 + lengthBytes == signature.Length;
    }

    public static byte[] ToRaw(ReadOnlySpan<byte> der, int coordinateLength)
    {
        if (!IsDer(der))
        {
            throw new CryptographicException("Signature is not DER encoded.");
        }

        var sequence = Asn1Sequence.GetInstance(Asn1Object.FromByteArray(der.ToArray()));
        if (sequence.Count != 2)
        {
            throw new CryptographicException("Invalid DER structure for GOST signature.");
        }

        var r = NormalizeCoordinate(((DerInteger)sequence[0]).PositiveValue.ToByteArrayUnsigned(), coordinateLength);
        var s = NormalizeCoordinate(((DerInteger)sequence[1]).PositiveValue.ToByteArrayUnsigned(), coordinateLength);

        var raw = new byte[coordinateLength * 2];
        s.CopyTo(raw.AsSpan(0, coordinateLength));
        r.CopyTo(raw.AsSpan(coordinateLength));
        return raw;
    }

    public static byte[] ToDer(ReadOnlySpan<byte> raw, int coordinateLength)
    {
        if (raw.Length != coordinateLength * 2)
        {
            throw new CryptographicException($"Raw GOST signature must be {coordinateLength * 2} bytes.");
        }

        var s = raw[..coordinateLength].ToArray();
        var r = raw[coordinateLength..].ToArray();

        var derSequence = new DerSequence(
            new DerInteger(new BigInteger(1, r)),
            new DerInteger(new BigInteger(1, s)));

        return derSequence.GetDerEncoded();
    }

    public static (BigInteger r, BigInteger s) DecodeComponents(ReadOnlySpan<byte> signature, int coordinateLength)
    {
        if (IsDer(signature))
        {
            var sequence = Asn1Sequence.GetInstance(Asn1Object.FromByteArray(signature.ToArray()));
            if (sequence.Count != 2)
            {
                throw new CryptographicException("Invalid DER structure for GOST signature.");
            }

            return (((DerInteger)sequence[0]).PositiveValue, ((DerInteger)sequence[1]).PositiveValue);
        }

        if (signature.Length == coordinateLength * 2)
        {
            var s = new byte[coordinateLength];
            var r = new byte[coordinateLength];
            signature[..coordinateLength].CopyTo(s);
            signature[coordinateLength..].CopyTo(r);
            return (new BigInteger(1, r), new BigInteger(1, s));
        }

        throw new CryptographicException("Signature payload is neither DER nor raw GOST format.");
    }

    private static byte[] NormalizeCoordinate(ReadOnlySpan<byte> value, int coordinateLength)
    {
        var trimmed = TrimLeadingZeros(value);
        if (trimmed.Length > coordinateLength)
        {
            throw new CryptographicException("Coordinate exceeds expected length.");
        }

        var output = new byte[coordinateLength];
        trimmed.CopyTo(output.AsSpan(coordinateLength - trimmed.Length));
        return output;
    }

    private static ReadOnlySpan<byte> TrimLeadingZeros(ReadOnlySpan<byte> value)
    {
        var index = 0;
        while (index < value.Length - 1 && value[index] == 0)
        {
            index++;
        }

        return value[index..];
    }
}