package crypto

import (
	"testing"
)

func TestEphemeralMarshalRoundTrip(t *testing.T) {
	params := DefaultParams(Security128)
	tags := []string{"word", "punct"}

	kpAlice, err := Generate(params, tags, testFactory)
	if err != nil {
		t.Fatalf("Generate(Alice): %v", err)
	}
	kpBob, err := Generate(params, tags, testFactory)
	if err != nil {
		t.Fatalf("Generate(Bob): %v", err)
	}

	msg, err := PrepareExchange(
		&kpAlice.Private, kpAlice.Public.Spore, kpBob.Public.Spore, params,
	)
	if err != nil {
		t.Fatalf("PrepareExchange: %v", err)
	}

	// Marshal.
	data, err := msg.MarshalEphemeral(params.N)
	if err != nil {
		t.Fatalf("MarshalEphemeral: %v", err)
	}

	if len(data) == 0 {
		t.Fatal("marshalled data should not be empty")
	}

	// Unmarshal.
	recovered, err := UnmarshalEphemeralMessage(data)
	if err != nil {
		t.Fatalf("UnmarshalEphemeralMessage: %v", err)
	}

	// Verify signature preserved.
	if recovered.Signature == nil {
		t.Fatal("recovered signature should not be nil")
	}
	// V2 stores only the first 16 bytes of the challenge (truncated binding).
	for i := range 16 {
		if recovered.Signature.Challenge[i] != msg.Signature.Challenge[i] {
			t.Errorf("challenge byte %d: got %02x, want %02x",
				i, recovered.Signature.Challenge[i], msg.Signature.Challenge[i])
		}
	}
	// Commitment is not preserved in V2 (recomputed by verifier).

	// Verify pattern length preserved.
	if len(recovered.Pattern) != len(msg.Pattern) {
		t.Errorf("pattern length = %d, want %d", len(recovered.Pattern), len(msg.Pattern))
	}

	// Count occupied sites in both.
	origOcc := 0
	for _, s := range msg.Pattern {
		if s.Occupied {
			origOcc++
		}
	}
	recOcc := 0
	for _, s := range recovered.Pattern {
		if s.Occupied {
			recOcc++
		}
	}
	if recOcc != origOcc {
		t.Errorf("occupied count = %d, want %d", recOcc, origOcc)
	}

	// Verify occupied sites have correct tags and projections.
	for i, s := range msg.Pattern {
		if !s.Occupied {
			continue
		}
		r := recovered.Pattern[i]
		if !r.Occupied {
			t.Errorf("site %d: should be occupied", i)
			continue
		}
		if r.TypeTag != s.TypeTag {
			t.Errorf("site %d: tag = %q, want %q", i, r.TypeTag, s.TypeTag)
		}
		if r.Projection != s.Projection {
			t.Errorf("site %d: projection = %d, want %d", i, r.Projection, s.Projection)
		}
		if r.Perm != s.Perm {
			t.Errorf("site %d: perm = %d, want %d", i, r.Perm, s.Perm)
		}
	}

	// Verify sender fingerprint PermDist preserved.
	for i := range 6 {
		if recovered.SenderFingerprint.PermDist[i] != msg.SenderFingerprint.PermDist[i] {
			t.Errorf("PermDist[%d] = %d, want %d",
				i, recovered.SenderFingerprint.PermDist[i], msg.SenderFingerprint.PermDist[i])
		}
	}
}

func TestEphemeralMarshalNil(t *testing.T) {
	var msg *EphemeralMessage
	_, err := msg.MarshalEphemeral(256)
	if err == nil {
		t.Error("expected error for nil message")
	}
}

func TestEphemeralUnmarshalTruncated(t *testing.T) {
	_, err := UnmarshalEphemeralMessage([]byte{0x01})
	if err == nil {
		t.Error("expected error for truncated data")
	}
}

func TestEphemeralUnmarshalBadVersion(t *testing.T) {
	data := make([]byte, 100)
	data[0] = 0xFF // bad version
	data[1] = 0xFF
	_, err := UnmarshalEphemeralMessage(data)
	if err == nil {
		t.Error("expected error for bad version")
	}
}

func TestEphemeralMarshalDeterministic(t *testing.T) {
	params := DefaultParams(Security128)
	tags := []string{"word", "punct"}

	kpAlice, _ := Generate(params, tags, testFactory)
	kpBob, _ := Generate(params, tags, testFactory)

	msg, _ := PrepareExchange(
		&kpAlice.Private, kpAlice.Public.Spore, kpBob.Public.Spore, params,
	)

	data1, _ := msg.MarshalEphemeral(params.N)
	data2, _ := msg.MarshalEphemeral(params.N)

	if len(data1) != len(data2) {
		t.Errorf("marshal lengths differ: %d vs %d", len(data1), len(data2))
	}
	for i := range data1 {
		if data1[i] != data2[i] {
			t.Errorf("marshal differs at byte %d: %02x vs %02x", i, data1[i], data2[i])
			break
		}
	}
}

func TestEphemeralMarshalPreservesSignatureIntegrity(t *testing.T) {
	// Verify that a marshalled+unmarshalled message preserves the
	// signature's internal consistency (challenge, commitment, response).
	// CompleteExchange uses the peer's known SporeFingerprint from their
	// spore (not from the wire message), so full exchange round-trip
	// verification is a protocol-level test, not a wire format test.
	params := DefaultParams(Security128)
	tags := []string{"word", "punct"}

	kpAlice, _ := Generate(params, tags, testFactory)
	kpBob, _ := Generate(params, tags, testFactory)

	msg, _ := PrepareExchange(
		&kpAlice.Private, kpAlice.Public.Spore, kpBob.Public.Spore, params,
	)

	data, err := msg.MarshalEphemeral(params.N)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	recovered, err := UnmarshalEphemeralMessage(data)
	if err != nil {
		t.Fatalf("unmarshal: %v", err)
	}

	// V2: truncated challenge (first 16 bytes) must be preserved.
	for i := range 16 {
		if recovered.Signature.Challenge[i] != msg.Signature.Challenge[i] {
			t.Errorf("challenge byte %d not preserved", i)
		}
	}
	// Commitment is not preserved in V2 format.

	// Response occupied count must match.
	origOcc := 0
	for _, s := range msg.Signature.Response {
		if s.Occupied {
			origOcc++
		}
	}
	recOcc := len(recovered.Signature.Response)
	if recOcc != origOcc {
		t.Errorf("signature response occupied = %d, want %d", recOcc, origOcc)
	}

	// V2 proof lengths must match response count.
	if len(recovered.Signature.Proof.LockIns) != recOcc {
		t.Errorf("proof lock-ins length = %d, want %d",
			len(recovered.Signature.Proof.LockIns), recOcc)
	}
}