// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mldsa_test

import (
	"fmt"
	"slices"
	"testing"

	"github.com/tink-crypto/tink-go/v2/insecuresecretdataaccess"
	"github.com/tink-crypto/tink-go/v2/internal/internalapi"
	"github.com/tink-crypto/tink-go/v2/internal/testing/wycheproof"
	"github.com/tink-crypto/tink-go/v2/keyset"
	"github.com/tink-crypto/tink-go/v2/secretdata"
	tinkmldsa "github.com/tink-crypto/tink-go/v2/signature/mldsa"
	"github.com/tink-crypto/tink-go/v2/signature"
	"github.com/tink-crypto/tink-go/v2/subtle/random"
	"github.com/tink-crypto/tink-go/v2/testutil"
)

func TestSignVerifyManager(t *testing.T) {
	message := random.GetRandomBytes(20)
	for _, tc := range []struct {
		name          string
		instance      tinkmldsa.Instance
		variant       tinkmldsa.Variant
		idRequirement uint32
	}{
		{
			name:          "TINK ML-DSA-44",
			instance:      tinkmldsa.MLDSA44,
			variant:       tinkmldsa.VariantTink,
			idRequirement: uint32(0x01020304),
		},
		{
			name:          "TINK ML-DSA-65",
			instance:      tinkmldsa.MLDSA65,
			variant:       tinkmldsa.VariantTink,
			idRequirement: uint32(0x01020304),
		},
		{
			name:          "TINK ML-DSA-87",
			instance:      tinkmldsa.MLDSA87,
			variant:       tinkmldsa.VariantTink,
			idRequirement: uint32(0x01020304),
		},
		{
			name:          "RAW ML-DSA-44",
			instance:      tinkmldsa.MLDSA44,
			variant:       tinkmldsa.VariantNoPrefix,
			idRequirement: uint32(0),
		},
		{
			name:          "RAW ML-DSA-65",
			instance:      tinkmldsa.MLDSA65,
			variant:       tinkmldsa.VariantNoPrefix,
			idRequirement: uint32(0),
		},
		{
			name:          "RAW ML-DSA-87",
			instance:      tinkmldsa.MLDSA87,
			variant:       tinkmldsa.VariantNoPrefix,
			idRequirement: uint32(0),
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			params, err := tinkmldsa.NewParameters(tc.instance, tc.variant)
			if err != nil {
				t.Fatalf("tinkmldsa.NewParameters(%v) err = %v, want nil", tc.variant, err)
			}
			publicKeyBytes, privateKeyBytes := getTestKeyPair(t, tc.instance)
			publicKey, err := tinkmldsa.NewPublicKey(publicKeyBytes, tc.idRequirement, params)
			if err != nil {
				t.Fatalf("tinkmldsa.NewPublicKey(%v, %v, %v) err = %v, want nil", publicKeyBytes, tc.idRequirement, params, err)
			}
			privateKey, err := tinkmldsa.NewPrivateKey(secretdata.NewBytesFromData(privateKeyBytes, insecuresecretdataaccess.Token{}), tc.idRequirement, params)
			if err != nil {
				t.Fatalf("tinkmldsa.NewPrivateKey(%v, %v, %v) err = %v, want nil", privateKeyBytes, tc.idRequirement, params, err)
			}

			// Signer verifier from keys.
			signer, err := tinkmldsa.NewSigner(privateKey, internalapi.Token{})
			if err != nil {
				t.Fatalf("tinkmldsa.NewSigner(%v, internalapi.Token{}) err = %v, want nil", privateKey, err)
			}
			sigBytes, err := signer.Sign(message)
			if err != nil {
				t.Fatalf("signer.Sign(%x) err = %v, want nil", message, err)
			}

			verifier, err := tinkmldsa.NewVerifier(publicKey, internalapi.Token{})
			if err != nil {
				t.Fatalf("tinkmldsa.NewVerifier(%v, internalapi.Token{}) err = %v, want nil", publicKey, err)
			}
			if err := verifier.Verify(sigBytes, message); err != nil {
				t.Errorf("verifier.Verify(%x, %x) err = %v, want nil", sigBytes, message, err)
			}

			// Signer verifier from keyset handle.
			km := keyset.NewManager()
			keyID, err := km.AddKey(privateKey)
			if err != nil {
				t.Fatalf("km.AddKey(%v) err = %v, want nil", privateKey, err)
			}
			if err := km.SetPrimary(keyID); err != nil {
				t.Fatalf("km.SetPrimary(%v) err = %v, want nil", keyID, err)
			}
			keysetHandle, err := km.Handle()
			if err != nil {
				t.Fatalf("km.Handle() err = %v, want nil", err)
			}
			publicKeysetHandle, err := keysetHandle.Public()
			if err != nil {
				t.Fatalf("keysetHandle.Public() err = %v, want nil", err)
			}

			signerFromKeyset, err := signature.NewSigner(keysetHandle)
			if err != nil {
				t.Fatalf("signature.NewSigner(%v) err = %v, want nil", keysetHandle, err)
			}
			sigBytesFromKeyset, err := signerFromKeyset.Sign(message)
			if err != nil {
				t.Fatalf("signerFromKeyset.Sign(%x) err = %v, want nil", message, err)
			}

			verifierFromKeyset, err := signature.NewVerifier(publicKeysetHandle)
			if err != nil {
				t.Fatalf("signature.NewVerifier() err = %v, want nil", err)
			}
			if err := verifierFromKeyset.Verify(sigBytesFromKeyset, message); err != nil {
				t.Errorf("verifierFromKeyset.Verify(%x, %x) err = %v, want nil", sigBytesFromKeyset, message, err)
			}
		})
	}
}

func TestVerifyFails(t *testing.T) {
	for _, tc := range []struct {
		name     string
		instance tinkmldsa.Instance
		variant  tinkmldsa.Variant
	}{
		{
			name:     "TINK ML-DSA-44",
			instance: tinkmldsa.MLDSA44,
			variant:  tinkmldsa.VariantTink,
		},
		{
			name:     "TINK ML-DSA-65",
			instance: tinkmldsa.MLDSA65,
			variant:  tinkmldsa.VariantTink,
		},
		{
			name:     "TINK ML-DSA-87",
			instance: tinkmldsa.MLDSA87,
			variant:  tinkmldsa.VariantTink,
		},
		{
			name:     "RAW ML-DSA-44",
			instance: tinkmldsa.MLDSA44,
			variant:  tinkmldsa.VariantNoPrefix,
		},
		{
			name:     "RAW ML-DSA-65",
			instance: tinkmldsa.MLDSA65,
			variant:  tinkmldsa.VariantNoPrefix,
		},
		{
			name:     "RAW ML-DSA-87",
			instance: tinkmldsa.MLDSA87,
			variant:  tinkmldsa.VariantNoPrefix,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			seed := random.GetRandomBytes(32)
			params, err := tinkmldsa.NewParameters(tc.instance, tc.variant)
			if err != nil {
				t.Fatalf("tinkmldsa.NewParameters(%v, %v) err = %v, want nil", tc.instance, tc.variant, err)
			}
			privateKey, err := tinkmldsa.NewPrivateKey(secretdata.NewBytesFromData(seed, insecuresecretdataaccess.Token{}), 0, params)
			if err != nil {
				t.Fatalf("tinkmldsa.NewPrivateKey(%v, %v, %v) err = %v, want nil", seed, 0, params, err)
			}
			publicKey, _ := privateKey.PublicKey()
			mldsaPublicKey, ok := publicKey.(*tinkmldsa.PublicKey)
			if !ok {
				t.Fatalf("privateKey.PublicKey() is not a *tinkmldsa.PublicKey")
			}
			signer, err := tinkmldsa.NewSigner(privateKey, internalapi.Token{})
			if err != nil {
				t.Fatalf("tinkmldsa.NewSigner(%v, internalapi.Token{}) err = %v, want nil", privateKey, err)
			}
			verifier, err := tinkmldsa.NewVerifier(mldsaPublicKey, internalapi.Token{})
			if err != nil {
				t.Fatalf("tinkmldsa.NewVerifier(%v, internalapi.Token{}) err = %v, want nil", mldsaPublicKey, err)
			}
			data := random.GetRandomBytes(20)
			signature, err := signer.Sign(data)
			if err != nil {
				t.Fatalf("signer.Sign(%x) err = %v, want nil", data, err)
			}

			prefix := signature[:len(mldsaPublicKey.OutputPrefix())]
			rawSignature := signature[len(mldsaPublicKey.OutputPrefix()):]

			// Modify the prefix.
			for i := 0; i < len(prefix); i++ {
				modifiedPrefix := slices.Clone(prefix)
				for j := 0; j < 8; j++ {
					modifiedPrefix[i] = byte(modifiedPrefix[i] ^ (1 << uint32(j)))
					s := slices.Concat(modifiedPrefix, rawSignature)
					if err := verifier.Verify(s, data); err == nil {
						t.Errorf("verifier.Verify(%x, data) err = nil, want error", s)
					}
				}
			}
			// Modify the signature.
			// This is slow, so we do not test every possible bit flip.
			for i := 0; i < len(rawSignature); i++ {
				modifiedRawSignature := slices.Clone(rawSignature)
				mask := byte(1 + random.GetRandomUint32()&0xFE)
				modifiedRawSignature[i] = byte(modifiedRawSignature[i] ^ mask)
				s := slices.Concat(prefix, modifiedRawSignature)
				if err := verifier.Verify(s, data); err == nil {
					t.Errorf("verifier.Verify(%x, data) err = nil, want error", s)
				}
			}
			// Modify the message.
			for i := 0; i < len(data); i++ {
				modifiedData := slices.Clone(data)
				for j := 0; j < 8; j++ {
					modifiedData[i] = byte(modifiedData[i] ^ (1 << uint32(j)))
					if err := verifier.Verify(signature, modifiedData); err == nil {
						t.Errorf("verifier.Verify(signature, %x) err = nil, want error", modifiedData)
					}
				}
			}
		})
	}
}

func TestSignVerifyCorrectness(t *testing.T) {
	for _, tc := range []struct {
		name     string
		instance tinkmldsa.Instance
		variant  tinkmldsa.Variant
	}{
		{
			name:     "TINK ML-DSA-44",
			instance: tinkmldsa.MLDSA44,
			variant:  tinkmldsa.VariantTink,
		},
		{
			name:     "TINK ML-DSA-65",
			instance: tinkmldsa.MLDSA65,
			variant:  tinkmldsa.VariantTink,
		},
		{
			name:     "TINK ML-DSA-87",
			instance: tinkmldsa.MLDSA87,
			variant:  tinkmldsa.VariantTink,
		},
		{
			name:     "RAW ML-DSA-44",
			instance: tinkmldsa.MLDSA44,
			variant:  tinkmldsa.VariantNoPrefix,
		},
		{
			name:     "RAW ML-DSA-65",
			instance: tinkmldsa.MLDSA65,
			variant:  tinkmldsa.VariantNoPrefix,
		},
		{
			name:     "RAW ML-DSA-87",
			instance: tinkmldsa.MLDSA87,
			variant:  tinkmldsa.VariantNoPrefix,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			seed := random.GetRandomBytes(32)
			params, err := tinkmldsa.NewParameters(tc.instance, tc.variant)
			if err != nil {
				t.Fatalf("tinkmldsa.NewParameters(%v, %v) err = %v, want nil", tc.instance, tc.variant, err)
			}
			privateKey, err := tinkmldsa.NewPrivateKey(secretdata.NewBytesFromData(seed, insecuresecretdataaccess.Token{}), 0, params)
			if err != nil {
				t.Fatalf("tinkmldsa.NewPrivateKey(%v, %v, %v) err = %v, want nil", seed, 0, params, err)
			}
			publicKey, _ := privateKey.PublicKey()
			mldsaPublicKey, ok := publicKey.(*tinkmldsa.PublicKey)
			if !ok {
				t.Fatalf("privateKey.PublicKey() is not a *tinkmldsa.PublicKey")
			}
			signer, err := tinkmldsa.NewSigner(privateKey, internalapi.Token{})
			if err != nil {
				t.Fatalf("tinkmldsa.NewSigner(%v, internalapi.Token{}) err = %v, want nil", privateKey, err)
			}
			verifier, err := tinkmldsa.NewVerifier(mldsaPublicKey, internalapi.Token{})
			if err != nil {
				t.Fatalf("tinkmldsa.NewVerifier(%v, internalapi.Token{}) err = %v, want nil", mldsaPublicKey, err)
			}
			for i := 0; i < 100; i++ {
				data := random.GetRandomBytes(20)
				signatureBytes, err := signer.Sign(data)
				if err != nil {
					t.Fatalf("signer.Sign(%x) err = %v, want nil", data, err)
				}
				if err := verifier.Verify(signatureBytes, data); err != nil {
					t.Errorf("verifier.Verify(%x, %x) err = %v, want nil", signatureBytes, data, err)
				}
			}
		})
	}
}

func TestWycheproof(t *testing.T) {
	vectors := []struct {
		filename string
		instance tinkmldsa.Instance
	}{
		{
			filename: "mldsa_44_verify_test.json",
			instance: tinkmldsa.MLDSA44,
		},
		{
			filename: "mldsa_65_verify_test.json",
			instance: tinkmldsa.MLDSA65,
		},
		{
			filename: "mldsa_87_verify_test.json",
			instance: tinkmldsa.MLDSA87,
		},
	}

	type wycheproofCase struct {
		wycheproof.Case
		Message   testutil.HexBytes `json:"msg"`
		Signature testutil.HexBytes `json:"sig"`
		Context   testutil.HexBytes `json:"ctx"`
	}

	type wycheproofGroup struct {
		wycheproof.Group
		PublicKey testutil.HexBytes `json:"publicKey"`
		Tests     []*wycheproofCase `json:"tests"`
	}

	type wycheproofSuite struct {
		wycheproof.Suite
		TestGroups []*wycheproofGroup `json:"testGroups"`
	}

	for _, v := range vectors {
		suite := new(wycheproofSuite)
		wycheproof.PopulateSuiteV1(t, suite, v.filename)

		for _, group := range suite.TestGroups {
			for _, test := range group.Tests {
				// The tink.Verifier interface does not support passing context.
				if len(test.Context) != 0 {
					continue
				}

				expectedPublicKeyParsingErr := slices.Contains(test.Flags, "IncorrectPublicKeyLength")

				caseName := fmt.Sprintf("%s:Case-%d", v.filename, test.CaseID)
				t.Run(caseName, func(t *testing.T) {
					params, err := tinkmldsa.NewParameters(v.instance, tinkmldsa.VariantNoPrefix)
					if err != nil {
						t.Fatalf("tinkmldsa.NewParameters(%v, tinkmldsa.VariantNoPrefix) err = %v, want nil", v.instance, err)
					}

					publicKey, err := tinkmldsa.NewPublicKey(group.PublicKey, 0, params)
					if expectedPublicKeyParsingErr {
						if err == nil {
							t.Fatalf("tinkmldsa.NewPublicKey(%x, 0, %+v) err = nil, want err", group.PublicKey, params)
						}
						// End subtest successfully.
						return
					}
					if err != nil {
						t.Fatalf("tinkmldsa.NewPublicKey(%x, 0, %+v) err = %v, want nil", group.PublicKey, params, err)
					}

					verifier, err := tinkmldsa.NewVerifier(publicKey, internalapi.Token{})
					if err != nil {
						t.Fatalf("tinkmldsa.NewVerifier(%v, internalapi.Token{}) err = %v, want nil", publicKey, err)
					}

					err = verifier.Verify(test.Signature, test.Message)
					switch test.Result {
					case "valid":
						if err != nil {
							t.Errorf("verifier.Verify() failed in a valid test case: %s", err)
						}
					case "invalid":
						if err == nil {
							t.Errorf("verifier.Verify() succeeded in an invalid test case")
						}
					default:
						t.Errorf("unsupported test result: %q", test.Result)
					}
				})
			}
		}
	}
}
