package auth

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestPasswordCipherDecryptsRSAOAEP256(t *testing.T) {
	t.Parallel()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	privatePEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}))

	cipher, err := NewPasswordCipher(privatePEM, "test-key")
	require.NoError(t, err)

	encrypted, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, &key.PublicKey, []byte("Admin@123"), nil)
	require.NoError(t, err)

	got, err := cipher.Decrypt(base64.StdEncoding.EncodeToString(encrypted), "test-key")
	require.NoError(t, err)
	require.Equal(t, "Admin@123", got)
}

func TestPasswordCipherRejectsWrongKeyID(t *testing.T) {
	t.Parallel()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	privatePEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}))

	cipher, err := NewPasswordCipher(privatePEM, "expected")
	require.NoError(t, err)

	_, err = cipher.Decrypt("bad", "other")
	require.Error(t, err)
	require.Contains(t, err.Error(), "key id mismatch")
}

func TestNewEphemeralPasswordCipherExposesPublicKey(t *testing.T) {
	t.Parallel()

	cipher, err := NewEphemeralPasswordCipher("ephemeral")
	require.NoError(t, err)

	publicKey := cipher.PublicKey()
	require.Equal(t, "ephemeral", publicKey.KeyID)
	require.Equal(t, "RSA-OAEP-256", publicKey.Algorithm)
	require.Contains(t, publicKey.PublicKey, "BEGIN PUBLIC KEY")
}

func TestPasswordCipherAcceptsFoldedPrivateKeyPEM(t *testing.T) {
	t.Parallel()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	privatePEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: mustMarshalPKCS8PrivateKey(t, key),
	}))
	foldedPEM := strings.Join(strings.Fields(privatePEM), " ")

	_, err = NewPasswordCipher(foldedPEM, "folded")
	require.NoError(t, err)
}

func TestPasswordCipherAcceptsEscapedNewlinePrivateKeyPEM(t *testing.T) {
	t.Parallel()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	privatePEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: mustMarshalPKCS8PrivateKey(t, key),
	}))
	escapedPEM := strings.ReplaceAll(strings.TrimSpace(privatePEM), "\n", `\n`)

	_, err = NewPasswordCipher(escapedPEM, "escaped")
	require.NoError(t, err)
}

func TestNormalizePrivateKeyPEMPreservesValidPEM(t *testing.T) {
	t.Parallel()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	privatePEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: mustMarshalPKCS8PrivateKey(t, key),
	}))

	require.Equal(t, strings.TrimSpace(privatePEM), strings.TrimSpace(normalizePrivateKeyPEM(privatePEM)))
}

func mustMarshalPKCS8PrivateKey(t *testing.T, key *rsa.PrivateKey) []byte {
	t.Helper()
	der, err := x509.MarshalPKCS8PrivateKey(key)
	require.NoError(t, err)
	return der
}