// Copyright (c) 2026 Alibaba Group and its affiliates

// 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 transparency

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/transparency-dev/merkle/rfc6962"

	"github.com/dashscope/dashscope-confidential-maas/pkg/utils"
)

var testRekorKey *ecdsa.PrivateKey
var testLogEntryKey *ecdsa.PrivateKey

func init() {
	var err error
	testRekorKey, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		panic("generate test Rekor key: " + err.Error())
	}
	testLogEntryKey, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		panic("generate test log entry key: " + err.Error())
	}
}

func marshalECDSAPublicKeyPEM(pub *ecdsa.PublicKey) string {
	der, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		panic("marshal test public key: " + err.Error())
	}
	return string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: der}))
}

func testServiceConfig() *RekorV1ServiceConfig {
	return &RekorV1ServiceConfig{
		RekorV1PubKeyPEM:  marshalECDSAPublicKeyPEM(&testRekorKey.PublicKey),
		LogEntryPubKeyPEM: marshalECDSAPublicKeyPEM(&testLogEntryKey.PublicKey),
	}
}

// makeVerification builds a valid RekorVerification (inclusion proof + checkpoint + SET)
// for a single-entry Merkle tree, signed with testRekorKey.
func makeVerification(entry *RekorEntry) *RekorVerification {
	raw, _ := base64.StdEncoding.DecodeString(entry.Body)
	rootBytes := rfc6962.DefaultHasher.HashLeaf(raw)
	rootHashHex := fmt.Sprintf("%x", rootBytes)
	rootHashB64 := base64.StdEncoding.EncodeToString(rootBytes)

	noteBody := fmt.Sprintf("rekor.sigstore.dev - 0\n1\n%s\n", rootHashB64)
	sigName := "rekor.sigstore.dev"
	noteHash := sha256.Sum256([]byte(noteBody))
	noteSig, _ := ecdsa.SignASN1(rand.Reader, testRekorKey, noteHash[:])

	pubKeyDER, _ := x509.MarshalPKIXPublicKey(&testRekorKey.PublicKey)
	keyHashFull := sha256.Sum256(pubKeyDER)
	var keyHashBytes [4]byte
	binary.BigEndian.PutUint32(keyHashBytes[:], binary.BigEndian.Uint32(keyHashFull[:4]))

	sigBytes := append(keyHashBytes[:], noteSig...)
	checkpoint := noteBody + "\n\u2014 " + sigName + " " + base64.StdEncoding.EncodeToString(sigBytes) + "\n"

	setPayload := struct {
		Body           string `json:"body"`
		IntegratedTime int64  `json:"integratedTime"`
		LogID          string `json:"logID"`
		LogIndex       int64  `json:"logIndex"`
	}{
		Body:           entry.Body,
		IntegratedTime: entry.IntegratedTime,
		LogID:          entry.LogID,
		LogIndex:       entry.LogIndex,
	}
	canonical, _ := json.Marshal(setPayload)
	setHash := sha256.Sum256(canonical)
	setSig, _ := ecdsa.SignASN1(rand.Reader, testRekorKey, setHash[:])

	return &RekorVerification{
		InclusionProof: &InclusionProof{
			TreeSize:   1,
			LogIndex:   0,
			RootHash:   rootHashHex,
			Hashes:     []string{},
			Checkpoint: checkpoint,
		},
		SignedEntryTimestamp: base64.StdEncoding.EncodeToString(setSig),
	}
}

func buildTestRekorEntry(t *testing.T, manifest *ReleaseManifest) *RekorEntry {
	t.Helper()

	jcs, err := utils.JCSCanonicalize(manifest)
	require.NoError(t, err)

	h := sha256.Sum256(jcs)
	pae := utils.DSSEPAE(DSSEPayloadType, jcs)
	ph := sha256.Sum256(pae)
	sig, err := ecdsa.SignASN1(rand.Reader, testLogEntryKey, ph[:])
	require.NoError(t, err)

	body := RekorBody{
		APIVersion: "0.0.1", Kind: "dsse",
		Spec: RekorBodySpec{
			PayloadHash: RekorPayloadHash{Algorithm: "sha256", Value: fmt.Sprintf("%x", h)},
			Signatures:  []RekorSignature{{Signature: base64.StdEncoding.EncodeToString(sig), Verifier: "mock"}},
		},
	}
	bb, err := json.Marshal(body)
	require.NoError(t, err)

	entry := &RekorEntry{
		Body: base64.StdEncoding.EncodeToString(bb), IntegratedTime: 1770726485, LogID: "test", LogIndex: 0,
	}
	entry.Verification = makeVerification(entry)
	return entry
}

func TestVerifyRekorV1PayloadHashMatch(t *testing.T) {
	manifest := &ReleaseManifest{
		SchemaVersion: 1,
		Measurements:  map[string]Measurement{"custom_metric": {Algorithm: "sha256", Value: "abc123"}},
	}

	rekorEntry := buildTestRekorEntry(t, manifest)
	entry := &ProofEntry{
		Type: string(LogTypeRekorV1),
		Data: &RekorV1Data{LogURL: "https://rekor.sigstore.dev", LogIndex: 0, RekorEntry: rekorEntry},
	}

	svc, err := newRekorV1Service(testServiceConfig())
	require.NoError(t, err)
	assert.NoError(t, svc.VerifyProofEntry(manifest, entry))
}

func TestVerifyRekorV1PayloadHashMismatch(t *testing.T) {
	manifest := &ReleaseManifest{
		SchemaVersion: 1,
		Measurements:  map[string]Measurement{"cvm_uki": {Algorithm: "sha384", Value: "abc123"}},
	}

	body := RekorBody{
		APIVersion: "0.0.1", Kind: "dsse",
		Spec: RekorBodySpec{
			PayloadHash: RekorPayloadHash{Algorithm: "sha256", Value: "0000000000000000000000000000000000000000000000000000000000000000"},
			Signatures:  []RekorSignature{{Signature: "mock", Verifier: "mock"}},
		},
	}
	bb, _ := json.Marshal(body)

	entry := &ProofEntry{
		Type: string(LogTypeRekorV1),
		Data: &RekorV1Data{
			LogURL: "https://rekor.sigstore.dev", LogIndex: 123,
			RekorEntry: &RekorEntry{Body: base64.StdEncoding.EncodeToString(bb), LogIndex: 123},
		},
	}

	svc, err := newRekorV1Service(testServiceConfig())
	require.NoError(t, err)
	err = svc.VerifyProofEntry(manifest, entry)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "payloadHash mismatch")
}

func TestVerifyRekorV1LogEntrySignature(t *testing.T) {
	manifest := &ReleaseManifest{
		SchemaVersion: 1,
		Measurements:  map[string]Measurement{"custom_metric": {Algorithm: "sha256", Value: "aaa"}},
	}

	jcs, err := utils.JCSCanonicalize(manifest)
	require.NoError(t, err)
	h := sha256.Sum256(jcs)
	pae := utils.DSSEPAE(DSSEPayloadType, jcs)
	ph := sha256.Sum256(pae)

	wrongKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err)
	wrongSig, err := ecdsa.SignASN1(rand.Reader, wrongKey, ph[:])
	require.NoError(t, err)

	body := RekorBody{
		APIVersion: "0.0.1", Kind: "dsse",
		Spec: RekorBodySpec{
			PayloadHash: RekorPayloadHash{Algorithm: "sha256", Value: fmt.Sprintf("%x", h)},
			Signatures:  []RekorSignature{{Signature: base64.StdEncoding.EncodeToString(wrongSig), Verifier: "mock"}},
		},
	}
	bb, _ := json.Marshal(body)

	bodyB64 := base64.StdEncoding.EncodeToString(bb)
	rekorEntry := &RekorEntry{Body: bodyB64, IntegratedTime: 1770726485, LogID: "test", LogIndex: 0}
	rekorEntry.Verification = makeVerification(rekorEntry)

	entry := &ProofEntry{
		Type: string(LogTypeRekorV1),
		Data: &RekorV1Data{LogURL: "https://rekor.sigstore.dev", LogIndex: 0, RekorEntry: rekorEntry},
	}

	svc, svcErr := newRekorV1Service(testServiceConfig())
	require.NoError(t, svcErr)
	err = svc.VerifyProofEntry(manifest, entry)
	assert.Error(t, err, "wrong key should be rejected")
	assert.Contains(t, err.Error(), "log entry signature verification failed")

	correctSig, err := ecdsa.SignASN1(rand.Reader, testLogEntryKey, ph[:])
	require.NoError(t, err)
	body.Spec.Signatures[0].Signature = base64.StdEncoding.EncodeToString(correctSig)
	bb, _ = json.Marshal(body)
	newRekorEntry := &RekorEntry{Body: base64.StdEncoding.EncodeToString(bb), IntegratedTime: 1770726485, LogID: "test", LogIndex: 0}
	newRekorEntry.Verification = makeVerification(newRekorEntry)
	entry.Data.(*RekorV1Data).RekorEntry = newRekorEntry

	assert.NoError(t, svc.VerifyProofEntry(manifest, entry), "correct key should pass")
}

func TestVerifyRekorV1WrongBodyKind(t *testing.T) {
	body := RekorBody{APIVersion: "0.0.1", Kind: "hashedrekord", Spec: RekorBodySpec{PayloadHash: RekorPayloadHash{Algorithm: "sha256", Value: "0000"}}}
	bb, _ := json.Marshal(body)

	entry := &ProofEntry{
		Type: string(LogTypeRekorV1),
		Data: &RekorV1Data{
			LogURL: "https://rekor.sigstore.dev", LogIndex: 123,
			RekorEntry: &RekorEntry{Body: base64.StdEncoding.EncodeToString(bb), LogIndex: 123},
		},
	}

	svc, svcErr := newRekorV1Service(testServiceConfig())
	require.NoError(t, svcErr)
	err := svc.VerifyProofEntry(&ReleaseManifest{SchemaVersion: 1, Measurements: map[string]Measurement{}}, entry)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "unexpected body kind")
}

func TestProofEntryJSON(t *testing.T) {
	rekorData := &RekorV1Data{
		LogURL: "https://rekor.sigstore.dev", LogIndex: 930102219,
		RekorEntry: &RekorEntry{
			Body: "base64body", IntegratedTime: 1770726485, LogID: "c0d23d6a", LogIndex: 930102219,
			Verification: &RekorVerification{
				InclusionProof: &InclusionProof{
					Checkpoint: "rekor.sigstore.dev - ...", Hashes: []string{"abc123", "def456"},
					LogIndex: 808197957, RootHash: "913379ae56b0", TreeSize: 808197961,
				},
				SignedEntryTimestamp: "MEQCIBU0Tt91...",
			},
		},
	}

	data, err := json.Marshal(ProofEntry{Type: string(LogTypeRekorV1), Data: rekorData})
	require.NoError(t, err)

	var decoded ProofEntry
	require.NoError(t, json.Unmarshal(data, &decoded))
	assert.Equal(t, "rekor-v1", decoded.Type)

	dataBytes, err := json.Marshal(decoded.Data)
	require.NoError(t, err)

	var rd RekorV1Data
	require.NoError(t, json.Unmarshal(dataBytes, &rd))
	assert.Equal(t, "https://rekor.sigstore.dev", rd.LogURL)
	assert.Equal(t, int64(930102219), rd.LogIndex)
	assert.NotNil(t, rd.RekorEntry)
}

func TestRekorBodyJSON(t *testing.T) {
	body := RekorBody{
		APIVersion: "0.0.1", Kind: "dsse",
		Spec: RekorBodySpec{
			PayloadHash: RekorPayloadHash{Algorithm: "sha256", Value: "9262eef4"},
			Signatures:  []RekorSignature{{Signature: "MEYCIQD...", Verifier: "LS0tLS1..."}},
		},
	}

	data, err := json.Marshal(body)
	require.NoError(t, err)

	var decoded RekorBody
	require.NoError(t, json.Unmarshal(data, &decoded))
	assert.Equal(t, "dsse", decoded.Kind)
	assert.Equal(t, "sha256", decoded.Spec.PayloadHash.Algorithm)
	assert.Len(t, decoded.Spec.Signatures, 1)
}

func TestHexDecode(t *testing.T) {
	cases := []struct {
		in      string
		want    []byte
		wantErr bool
	}{
		{"", []byte{}, false},
		{"00", []byte{0}, false},
		{"ff", []byte{255}, false},
		{"0123456789abcdef", []byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef}, false},
		{"0", nil, true},
		{"zz", nil, true},
	}
	for _, c := range cases {
		got, err := hex.DecodeString(c.in)
		if c.wantErr {
			assert.Error(t, err, "hex.DecodeString(%q) should fail", c.in)
		} else {
			assert.NoError(t, err, "hex.DecodeString(%q) should succeed", c.in)
			assert.Equal(t, c.want, got)
		}
	}
}

func TestVerifyPayloadHash(t *testing.T) {
	manifest := &ReleaseManifest{
		SchemaVersion: 1,
		Measurements:  map[string]Measurement{"test": {Algorithm: "sha256", Value: "abc"}},
	}
	jcs, err := utils.JCSCanonicalize(manifest)
	require.NoError(t, err)
	h := sha256.Sum256(jcs)

	assert.NoError(t, verifyPayloadHash(manifest, fmt.Sprintf("%x", h)))
	assert.Error(t, verifyPayloadHash(manifest, "0000000000000000000000000000000000000000000000000000000000000000"))
}
