package crypto

import (
	"bytes"
	"io"
	"strings"
	"testing"

	"github.com/ProtonMail/go-crypto/openpgp"
	"github.com/ProtonMail/go-crypto/openpgp/armor"
)

// genTestKey creates an OpenPGP keypair, returns the armored public key (as the
// service would export from Vault) and an unwrap func that decrypts the wrapped
// DEK with the private key — simulating the Vault GPG engine's decrypt endpoint
// (which returns the plaintext of a whole OpenPGP message).
func genTestKey(t *testing.T) (armoredPub []byte, unwrap UnwrapFunc) {
	t.Helper()
	ent, err := openpgp.NewEntity("logarchiver-test", "unit test", "test@unkin.net", nil)
	if err != nil {
		t.Fatalf("NewEntity: %v", err)
	}
	var buf bytes.Buffer
	w, err := armor.Encode(&buf, openpgp.PublicKeyType, nil)
	if err != nil {
		t.Fatalf("armor encode: %v", err)
	}
	if err := ent.Serialize(w); err != nil {
		t.Fatalf("serialize public: %v", err)
	}
	_ = w.Close()

	unwrap = func(wrapped []byte) ([]byte, error) {
		md, err := openpgp.ReadMessage(bytes.NewReader(wrapped), openpgp.EntityList{ent}, nil, nil)
		if err != nil {
			return nil, err
		}
		return io.ReadAll(md.UnverifiedBody)
	}
	return buf.Bytes(), unwrap
}

func TestRoundTrip(t *testing.T) {
	armoredPub, unwrap := genTestKey(t)
	pub, err := LoadPublicKey(armoredPub)
	if err != nil {
		t.Fatalf("LoadPublicKey: %v", err)
	}
	if len(pub.Fingerprint) != 40 {
		t.Errorf("fingerprint = %q, want 40 hex chars", pub.Fingerprint)
	}
	if pub.Fingerprint != strings.ToUpper(pub.Fingerprint) {
		t.Errorf("fingerprint should be uppercase: %q", pub.Fingerprint)
	}

	// A multi-line NDJSON payload larger than the frame size (forces >1 frame).
	var payload bytes.Buffer
	for i := 0; i < 5000; i++ {
		payload.WriteString(`{"host":"node-1","message":"line `)
		payload.WriteString(strings.Repeat("x", 50))
		payload.WriteString(`"}` + "\n")
	}
	plaintext := payload.Bytes()

	var sealed bytes.Buffer
	res, err := Seal(&sealed, plaintext, pub, "logarchive", 4096)
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}
	if res.RawBytes != int64(len(plaintext)) {
		t.Errorf("RawBytes = %d, want %d", res.RawBytes, len(plaintext))
	}
	if int64(sealed.Len()) != res.StoredBytes {
		t.Errorf("StoredBytes = %d, buffer = %d", res.StoredBytes, sealed.Len())
	}
	// Compression should shrink this highly repetitive payload.
	if res.StoredBytes >= res.RawBytes {
		t.Errorf("stored (%d) not smaller than raw (%d)", res.StoredBytes, res.RawBytes)
	}
	// Design property: only a tiny wrapped DEK goes to Vault, regardless of size.
	if res.Header.WrappedDEKLen > 4096 {
		t.Errorf("wrapped DEK unexpectedly large: %d bytes", res.Header.WrappedDEKLen)
	}
	if res.Header.KeyFingerprint != pub.Fingerprint {
		t.Errorf("header fingerprint mismatch")
	}

	var out bytes.Buffer
	if err := Open(bytes.NewReader(sealed.Bytes()), &out, unwrap); err != nil {
		t.Fatalf("Open: %v", err)
	}
	if !bytes.Equal(out.Bytes(), plaintext) {
		t.Fatalf("round-trip mismatch: got %d bytes, want %d", out.Len(), len(plaintext))
	}
}

func TestRoundTripEmpty(t *testing.T) {
	armoredPub, unwrap := genTestKey(t)
	pub, _ := LoadPublicKey(armoredPub)
	var sealed bytes.Buffer
	if _, err := Seal(&sealed, []byte{}, pub, "k", 4096); err != nil {
		t.Fatalf("Seal empty: %v", err)
	}
	var out bytes.Buffer
	if err := Open(bytes.NewReader(sealed.Bytes()), &out, unwrap); err != nil {
		t.Fatalf("Open empty: %v", err)
	}
	if out.Len() != 0 {
		t.Errorf("empty round-trip produced %d bytes", out.Len())
	}
}

func TestTamperDetected(t *testing.T) {
	armoredPub, unwrap := genTestKey(t)
	pub, _ := LoadPublicKey(armoredPub)
	var sealed bytes.Buffer
	if _, err := Seal(&sealed, []byte("hello world\n"), pub, "k", 4096); err != nil {
		t.Fatalf("Seal: %v", err)
	}
	data := sealed.Bytes()
	// Flip a byte near the end (inside a frame's ciphertext/tag).
	data[len(data)-3] ^= 0xff
	var out bytes.Buffer
	if err := Open(bytes.NewReader(data), &out, unwrap); err == nil {
		t.Fatalf("expected GCM authentication failure on tampered ciphertext")
	}
}

func TestBadMagic(t *testing.T) {
	_, _, err := ReadHeader(bytes.NewReader([]byte("NOTLARC.....")))
	if err == nil {
		t.Fatalf("expected bad-magic error")
	}
}

func TestReadHeaderFields(t *testing.T) {
	armoredPub, _ := genTestKey(t)
	pub, _ := LoadPublicKey(armoredPub)
	var sealed bytes.Buffer
	if _, err := Seal(&sealed, []byte("x\n"), pub, "logarchive", 4096); err != nil {
		t.Fatalf("Seal: %v", err)
	}
	hdr, wrapped, err := ReadHeader(bytes.NewReader(sealed.Bytes()))
	if err != nil {
		t.Fatalf("ReadHeader: %v", err)
	}
	if hdr.KeyName != "logarchive" {
		t.Errorf("KeyName = %q", hdr.KeyName)
	}
	if hdr.Compression != "zstd" || hdr.Cipher != "AES-256-GCM" {
		t.Errorf("algo metadata wrong: %+v", hdr)
	}
	if len(wrapped) != hdr.WrappedDEKLen {
		t.Errorf("wrapped len %d != header %d", len(wrapped), hdr.WrappedDEKLen)
	}
}

func TestDigestArmoredStable(t *testing.T) {
	// Guards the test helper used elsewhere; identical input -> identical digest.
	a := digestArmored([]byte("abc"))
	b := digestArmored([]byte("abc"))
	if a != b || a == "" {
		t.Errorf("digestArmored not stable: %q %q", a, b)
	}
}