package smime

import (
	"bytes"
	"crypto"
	"crypto/x509"
	"strings"
	"testing"
)

func testIdentity(t *testing.T, email string) (*x509.Certificate, crypto.PrivateKey) {
	t.Helper()
	certPEM, keyPEM, err := GenerateSelfSigned(email, DefaultValidity)
	if err != nil {
		t.Fatalf("GenerateSelfSigned: %v", err)
	}
	cert, err := ParseCertPEM(certPEM)
	if err != nil {
		t.Fatalf("ParseCertPEM: %v", err)
	}
	key, err := ParseKeyPEM(keyPEM)
	if err != nil {
		t.Fatalf("ParseKeyPEM: %v", err)
	}
	return cert, key
}

// testEntity uses CRLF line endings already, since Sign/Encrypt canonicalize the
// body to CRLF (MIME's wire form) before signing/encrypting — a round trip through
// either normalizes bare LF to CRLF, so tests compare against the canonical form.
func testEntity() Entity {
	return Entity{
		Headers: []string{"Content-Type: text/plain; charset=utf-8"},
		Body:    []byte("hello world\r\nsecond line\r\n"),
	}
}

func TestSignVerifyRoundTrip(t *testing.T) {
	cert, key := testIdentity(t, "alice@example.com")
	orig := testEntity()

	signed, err := Sign(orig, cert, key)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}
	if ct := HeaderValue(signed.Headers, "Content-Type"); !strings.HasPrefix(ct, "multipart/signed") {
		t.Fatalf("unexpected Content-Type: %q", ct)
	}

	inner, signer, err := VerifySigned(signed)
	if err != nil {
		t.Fatalf("VerifySigned: %v", err)
	}
	if signer == nil || signer.Subject.CommonName != "alice@example.com" {
		t.Fatalf("unexpected signer: %+v", signer)
	}
	if !bytes.Equal(inner.Body, orig.Body) {
		t.Fatalf("body mismatch: got %q want %q", inner.Body, orig.Body)
	}
	if HeaderValue(inner.Headers, "Content-Type") != HeaderValue(orig.Headers, "Content-Type") {
		t.Fatalf("header mismatch: got %v want %v", inner.Headers, orig.Headers)
	}
}

func TestVerifySignedDetectsTampering(t *testing.T) {
	cert, key := testIdentity(t, "alice@example.com")
	signed, err := Sign(testEntity(), cert, key)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	tampered := string(signed.Body)
	tampered = strings.Replace(tampered, "hello world", "hello WORLD", 1)
	signed.Body = []byte(tampered)

	inner, signer, err := VerifySigned(signed)
	if err == nil {
		t.Fatal("expected verification error for tampered content, got nil")
	}
	if signer != nil {
		t.Fatalf("expected nil signer on failed verification, got %+v", signer)
	}
	// The tampered body should still come back for display purposes even though
	// verification failed.
	if !bytes.Contains(inner.Body, []byte("hello WORLD")) {
		t.Fatalf("expected tampered body returned alongside the error, got %q", inner.Body)
	}
}

func TestVerifySignedWrongSignerCert(t *testing.T) {
	cert, key := testIdentity(t, "alice@example.com")
	other, _ := testIdentity(t, "mallory@example.com")

	signed, err := Sign(testEntity(), cert, key)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}
	_, signer, err := VerifySigned(signed)
	if err != nil {
		t.Fatalf("VerifySigned: %v", err)
	}
	if signer.Subject.CommonName == other.Subject.CommonName {
		t.Fatal("signer should not match an unrelated certificate")
	}
}

func TestEncryptDecryptRoundTrip(t *testing.T) {
	cert, key := testIdentity(t, "bob@example.com")
	orig := testEntity()

	encrypted, err := Encrypt(orig, []*x509.Certificate{cert})
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}
	if ct := HeaderValue(encrypted.Headers, "Content-Type"); !strings.HasPrefix(ct, "application/pkcs7-mime") {
		t.Fatalf("unexpected Content-Type: %q", ct)
	}
	if bytes.Contains(encrypted.Body, orig.Body) {
		t.Fatal("encrypted body should not contain the plaintext")
	}

	decrypted, err := Decrypt(encrypted, cert, key)
	if err != nil {
		t.Fatalf("Decrypt: %v", err)
	}
	if !bytes.Equal(decrypted.Body, orig.Body) {
		t.Fatalf("body mismatch: got %q want %q", decrypted.Body, orig.Body)
	}
	if HeaderValue(decrypted.Headers, "Content-Type") != HeaderValue(orig.Headers, "Content-Type") {
		t.Fatalf("header mismatch: got %v want %v", decrypted.Headers, orig.Headers)
	}
}

func TestDecryptWrongKeyFails(t *testing.T) {
	cert, _ := testIdentity(t, "bob@example.com")
	otherCert, otherKey := testIdentity(t, "mallory@example.com")

	encrypted, err := Encrypt(testEntity(), []*x509.Certificate{cert})
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}
	if _, err := Decrypt(encrypted, otherCert, otherKey); err == nil {
		t.Fatal("expected decryption with the wrong key to fail")
	}
}

func TestEncryptMultipleRecipientsBothCanDecrypt(t *testing.T) {
	senderCert, senderKey := testIdentity(t, "alice@example.com")
	recipCert, recipKey := testIdentity(t, "bob@example.com")
	orig := testEntity()

	encrypted, err := Encrypt(orig, []*x509.Certificate{senderCert, recipCert})
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}

	senderCopy, err := Decrypt(encrypted, senderCert, senderKey)
	if err != nil {
		t.Fatalf("sender Decrypt: %v", err)
	}
	if !bytes.Equal(senderCopy.Body, orig.Body) {
		t.Fatal("sender's own copy did not decrypt to the original body")
	}

	recipCopy, err := Decrypt(encrypted, recipCert, recipKey)
	if err != nil {
		t.Fatalf("recipient Decrypt: %v", err)
	}
	if !bytes.Equal(recipCopy.Body, orig.Body) {
		t.Fatal("recipient's copy did not decrypt to the original body")
	}
}

// TestSignThenEncryptNestedRoundTrip covers the "sign and encrypt" compose option:
// the plaintext is signed, then the whole signed entity is encrypted (opaque
// nesting), matching how webui's compose handler applies both transforms together.
func TestSignThenEncryptNestedRoundTrip(t *testing.T) {
	senderCert, senderKey := testIdentity(t, "alice@example.com")
	recipCert, recipKey := testIdentity(t, "bob@example.com")
	orig := testEntity()

	signed, err := Sign(orig, senderCert, senderKey)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}
	encrypted, err := Encrypt(signed, []*x509.Certificate{recipCert})
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}

	decrypted, err := Decrypt(encrypted, recipCert, recipKey)
	if err != nil {
		t.Fatalf("Decrypt: %v", err)
	}
	if !strings.HasPrefix(HeaderValue(decrypted.Headers, "Content-Type"), "multipart/signed") {
		t.Fatalf("expected the decrypted layer to still be multipart/signed, got %q", HeaderValue(decrypted.Headers, "Content-Type"))
	}

	inner, signer, err := VerifySigned(decrypted)
	if err != nil {
		t.Fatalf("VerifySigned: %v", err)
	}
	if signer.Subject.CommonName != "alice@example.com" {
		t.Fatalf("unexpected signer: %+v", signer)
	}
	if !bytes.Equal(inner.Body, orig.Body) {
		t.Fatalf("body mismatch after unwrapping both layers: got %q want %q", inner.Body, orig.Body)
	}
}

func TestImportPKCS12RejectsBadPassword(t *testing.T) {
	if _, _, err := ImportPKCS12([]byte("not a real pkcs12 file"), "whatever"); err == nil {
		t.Fatal("expected an error decoding garbage PKCS#12 data")
	}
}