// Copyright 2025 Sun Yimin. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

//go:build go1.24

package mldsa

import (
	"bytes"
	"encoding/hex"
	"testing"
)

// https://github.com/usnistgov/ACVP-Server/blob/master/gen-val/json-files/ML-DSA-keyGen-FIPS204/internalProjection.json
var kenGen87InternalProjectionCases = []struct {
	seed string
	pk   string
	sk   string
}{
	{
		"631AFC2A36A57E1D090DADC2791D486D72C69A9AABF97990C573214846FE5B64",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
	},
	{
		"939B3D7BCA29F54EB75901AFE54794BACDCCA65BEA120AD522BD95C38466BB84",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
	},
	{
		"157E86FE1BDF56B2E4FA121D625010DFBF5C79DA3C507645CFF5063D640A2ED4",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
	},
}

func TestKeyGen87(t *testing.T) {
	for _, c := range kenGen87InternalProjectionCases {
		xi, _ := hex.DecodeString(c.seed)
		pk, _ := hex.DecodeString(c.pk)
		sk, _ := hex.DecodeString(c.sk)
		priv, err := NewKey87(xi)
		if err != nil {
			t.Fatalf("NewPrivateKey65 failed: %v", err)
		}
		pub := priv.PublicKey()
		pubBytes := pub.Bytes()
		if !bytes.Equal(pubBytes, pk) {
			t.Errorf("Public key mismatch: got %x, want %x", pubBytes, pk)
		}
		pub2, err := NewPublicKey87(pubBytes)
		if err != nil {
			t.Fatalf("faile to parse public key: %v", err)
		}
		if !pub.Equal(pub2) {
			t.Errorf("Public key not equal: got %x, want %x", pubBytes, pub2.Bytes())
		}
		privBytes := priv.PrivateKey87.Bytes()
		if !bytes.Equal(privBytes, sk) {
			t.Errorf("Private key mismatch: got %x, want %x", privBytes, sk)
		}

		priv2, err := NewPrivateKey87(privBytes)
		if err != nil {
			t.Fatalf("failed to parse private key: %v", err)
		}
		if !priv.Equal(priv2) {
			t.Errorf("Private key not equal: got %x, want %x", privBytes, priv2.Bytes())
		}
	}
}

// https://raw.githubusercontent.com/usnistgov/ACVP-Server/refs/heads/master/gen-val/json-files/ML-DSA-sigGen-FIPS204/internalProjection.json
var sigGen87InternalProjectionCases = []struct {
	mu  string
	pk  string
	sk  string
	sig string
}{
	{
		"8D47E774CE2E855173438A096E8795B639A36FE254F5A66A9A46E934B48F544512347D7EE02E9439AA6406436336FA9CF7EB647760DF720AC5C8A363672FB579",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 9254 characters]",
	},
	{
		"EC3282A9CF37A26EE42DC78D0F03EFE0AFA889A7526984C1030F6D3F6E0D7D9F303C8E1AB4D23BBB36D3D409C57CAAF2267C7AC427671E7B4D367C0AF3D942B9",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 9254 characters]",
	},
	{
		"C92B3682759601926479B5B92656A5A4832F692973162D6555645F1F8A61EFC67869762779EC56A142668D811147DA2F0A033DE523C17E4053B06FCB3744F8F2",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 9254 characters]",
	},
	{
		"5F85AB235F88D1300D149A3E053224256DCFAB16BA5601F42AE3F1F300E29776A366610DF5D6BBAE9015A698BD50EB0F32E3F74E7C8913C193A1F91B376146B4",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 9254 characters]",
	},
	{
		"5517EAAEA8A2B37F957954A4464C0F361916E1D85C0A769B849DFF42074EF656275BE1238860FAF6CEA1302569FACC201699A06320F189363B730EA2998221C5",
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 9254 characters]",
	},
}

func TestSign87(t *testing.T) {
	var seed [32]byte
	for _, c := range sigGen87InternalProjectionCases {
		sk, _ := hex.DecodeString(c.sk)
		pk, _ := hex.DecodeString(c.pk)
		sig, _ := hex.DecodeString(c.sig)
		mu, _ := hex.DecodeString(c.mu)
		priv, err := NewPrivateKey87(sk)
		if err != nil {
			t.Fatalf("NewPrivateKey87 failed: %v", err)
		}
		sig2, err := priv.signInternal(seed[:], mu)
		if err != nil {
			t.Fatalf("failed to sign: %v", err)
		}
		if !bytes.Equal(sig2[:lambda256/4], sig[:lambda256/4]) {
			t.Errorf("signature mismatch(signer's commitment hash): got %x, want %x", sig2[:lambda256/4], sig[:lambda256/4])
		}
		if !bytes.Equal(sig2[lambda256/4:lambda256/4+encodingSize20*l87], sig[lambda256/4:lambda256/4+encodingSize20*l87]) {
			t.Errorf("signature mismatch(signer's response): got %x, want %x", sig2[lambda256/4:lambda256/4+encodingSize20*l87], sig[lambda256/4:lambda256/4+encodingSize20*l87])
		}
		if !bytes.Equal(sig2[lambda256/4+encodingSize20*l87:], sig[lambda256/4+encodingSize20*l87:]) {
			t.Errorf("signature mismatch(hint): got %x, want %x", sig2[lambda256/4+encodingSize20*l87:], sig[lambda256/4+encodingSize20*l87:])
		}
		pub, err := NewPublicKey87(pk)
		if err != nil {
			t.Fatalf("NewPublicKey87 failed: %v", err)
		}
		if !pub.verifyInternal(sig, mu) {
			t.Error("signature verification failed")
		}
	}
}

// https://raw.githubusercontent.com/usnistgov/ACVP-Server/refs/heads/master/gen-val/json-files/ML-DSA-sigVer-FIPS204/internalProjection.json
var sigVer87InternalProjectionCases = []struct {
	pk      string
	sk      string
	message string
	context string
	sig     string
	passed  bool
}{
	{
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 10530 characters]",
		"4318C3522B195DFD49B792F71AB5D8CAE4C8292797CAAC66579AACE3AC88F741A1A80AD29E6C05A8C3A134D57D3EA500544E33F8AA99B84DEA62BF31C2209E7D8D1036C0B62C1818C0D4C14F1797C916CB6E044B068552D768C0FEB785C67D84F14B6C4755716E59993A7DFD71210DE91A44F83FA58FC9F9588B312486B2523ED132712580E35BC4ADAADBD07FAA6399E58469D7054EB76271ED2AB0351C03740AFCA7D4BBF36E831BFE8D5CF23AE7833C664189275C3B78062D18EF6777C23B07CD0B662DA000260E903441C7BD908ADC9228F7580A421563D574BC1A7C",
		"[encoded data, 9254 characters]",
		false,
	},
	{
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 13774 characters]",
		"",
		"[encoded data, 9254 characters]",
		false,
	},
	{
		"[encoded data, 5184 characters]",
		"[encoded data, 9792 characters]",
		"[encoded data, 4290 characters]",
		"B3C4558FCBA8C92EC456A5CE50151142F5176E102BC843DE01EA64FB93465F50F0E7A9A6FC930A413EC5AE5699168AFE75C6FBD0C655085C6A5222D6E908F472AD3FE5DA148709F0E835653325AE3123C04D0387B80D8FD5C1BD2D0683F276A4F9D305CB3A4A4C794EB3DD562BC91A9859A6C860ACED45B30A1F25600057A4C68C90FC40468ACA89AEB074FEC09AF5E61E2D5E27B90335F6134C5504D4FA29C2B29EF120AF39C5B8CE8164E2E359FF6023F95B291A56647EB1AC790C42980B12925858C0FB19DF7A56797E1E82E7A38D",
		"[encoded data, 9254 characters]",
		true,
	},
}

func TestVerify87(t *testing.T) {
	for _, c := range sigVer87InternalProjectionCases {
		//sk, _ := hex.DecodeString(c.sk)
		pk, _ := hex.DecodeString(c.pk)
		sig, _ := hex.DecodeString(c.sig)
		msg, _ := hex.DecodeString(c.message)
		ctx, _ := hex.DecodeString(c.context)

		pub, err := NewPublicKey87(pk)
		if err != nil {
			t.Fatalf("NewPublicKey87 failed: %v", err)
		}
		if pub.Verify(sig, msg, ctx) != c.passed {
			t.Errorf("Verify failed")
		}
	}
}
