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

import (
	"encoding/binary"
	"slices"
	"testing"

	"crypto/rand"
)

func TestNewAddress(t *testing.T) {
	a := newAddress()
	want := [32]byte{}
	if !slices.Equal(a[:], want[:]) {
		t.Errorf("newAddress() = %v, want %v", a[:], want[:])
	}
}

func TestAddressCopy(t *testing.T) {
	a := newAddress()
	a.setTypeAndClear(addressWOTSHash)
	a.setLayerAddress(0xAB)
	a.setTreeAddress(0x0123456789ABCDEF)
	a.setKeyPairAddress(0xDEADBEEF)
	a.setChainAddress(0xC0FEBABE)
	a.setHashAddress(0xF00DD00B)
	a2 := a.copy()
	if !slices.Equal(a[:], a2[:]) {
		t.Errorf("a.copy() = %v, want %v", a2[:], a[:])
	}
}

func TestAddressSetLayerAddress(t *testing.T) {
	a := newAddress()
	a.setLayerAddress(0xDEADBEEF)
	want := address{0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	if !slices.Equal(a[:], want[:]) {
		t.Errorf("a.setLayerAddress(0xDEADBEEF) = %v, want %v", a[:], want[:])
	}
}

func TestAddressSetTreeAddress(t *testing.T) {
	a := newAddress()
	a.setTreeAddress(0xF123456789ABCDEF)
	want := address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xF1, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	if !slices.Equal(a[:], want[:]) {
		t.Errorf("a.setTreeAddress(0xF123456789ABCDEF) = %v, want %v", a[:], want[:])
	}
}

func TestAddressSetTypeAndClear(t *testing.T) {
	for _, tc := range []struct {
		name     string
		addrType addressType
		want     address
	}{
		{
			name:     "adrsWotsHash",
			addrType: addressWOTSHash,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
		{
			name:     "adrsWotsPk",
			addrType: addressWOTSPk,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
		{
			name:     "adrsTree",
			addrType: addressTree,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
		{
			name:     "adrsForsTree",
			addrType: addressFORSTree,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
		{
			name:     "adrsForsRoots",
			addrType: addressFORSRoots,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
		{
			name:     "adrsWotsPrf",
			addrType: addressWOTSPrf,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
		{
			name:     "adrsForsPrf",
			addrType: addressFORSPrf,
			want:     address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		},
	} {
		a := newAddress()
		rand.Read(a[20:32])
		a.setTypeAndClear(tc.addrType)
		if !slices.Equal(a[:], tc.want[:]) {
			t.Fatalf("a.setTypeAndClear(%v) = %v, want %v", tc.name, a[:], tc.want[:])
		}
		if binary.BigEndian.Uint32(a[16:20]) != uint32(tc.addrType) {
			t.Errorf("a[16:20] = %v, want %v", binary.BigEndian.Uint32(a[16:20]), uint32(tc.addrType))
		}
	}
}

func TestAddressSetGetKeyPairAddress(t *testing.T) {
	a := newAddress()
	a.setKeyPairAddress(0xDEADBEEF)
	want := address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	if !slices.Equal(a[:], want[:]) {
		t.Fatalf("a.setKeyPairAddress(0xDEADBEEF) = %v, want %v", a[:], want[:])
	}
	if a.keyPairAddress() != 0xDEADBEEF {
		t.Errorf("a.getKeyPairAddress() = %v, want 0xDEADBEEF", a.keyPairAddress())
	}
}

func TestAddressSetChainAddress(t *testing.T) {
	a := newAddress()
	a.setChainAddress(0xDEADBEEF)
	want := address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x00, 0x00, 0x00}
	if !slices.Equal(a[:], want[:]) {
		t.Errorf("a.setChainAddress(0xDEADBEEF) = %v, want %v", a[:], want[:])
	}
}

func TestAddressSetTreeHeight(t *testing.T) {
	a := newAddress()
	a.setTreeHeight(0xDEADBEEF)
	want := address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x00, 0x00, 0x00}
	if !slices.Equal(a[:], want[:]) {
		t.Errorf("a.setTreeHeight(0xDEADBEEF) = %v, want %v", a[:], want[:])
	}
}

func TestAddressSetHashAddress(t *testing.T) {
	a := newAddress()
	a.setHashAddress(0xDEADBEEF)
	want := address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xDE, 0xAD, 0xBE, 0xEF}
	if !slices.Equal(a[:], want[:]) {
		t.Errorf("a.setHashAddress(0xDEADBEEF) = %v, want %v", a[:], want[:])
	}
}

func TestAddressSetGetTreeIndex(t *testing.T) {
	a := newAddress()
	a.setTreeIndex(0xDEADBEEF)
	want := address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xDE, 0xAD, 0xBE, 0xEF}
	if !slices.Equal(a[:], want[:]) {
		t.Fatalf("a.setTreeIndex(0xDEADBEEF) = %v, want %v", a[:], want[:])
	}
	if a.treeIndex() != 0xDEADBEEF {
		t.Errorf("a.getTreeIndex() = %v, want 0xDEADBEEF", a.treeIndex())
	}
}

func TestAddressCompress(t *testing.T) {
	a := newAddress()
	a.setTypeAndClear(addressWOTSHash)
	a.setLayerAddress(0xAB)
	a.setTreeAddress(0x0123456789ABCDEF)
	a.setKeyPairAddress(0xDEADBEEF)
	a.setChainAddress(0xC0FEBABE)
	a.setHashAddress(0xF00DD00B)
	want := []byte{0xAB, 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0x00, 0xDE, 0xAD, 0xBE, 0xEF, 0xC0, 0xFE, 0xBA, 0xBE, 0xF0, 0x0D, 0xD0, 0x0B}
	if !slices.Equal(a.compress(), want) {
		t.Errorf("a.compress() = %v, want %v", a.compress(), want)
	}
}
