package auth

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"strings"
	"sync"
)

const DefaultPasswordCipherKeyID = "login-password-rsa-oaep-v1"

type PasswordCipher struct {
	mu         sync.RWMutex
	privateKey *rsa.PrivateKey
	keyID      string
	publicPEM  string
}

type PasswordCipherPublicKey struct {
	KeyID     string `json:"key_id"`
	Algorithm string `json:"algorithm"`
	PublicKey string `json:"public_key"`
}

func NewPasswordCipher(privateKeyPEM, keyID string) (*PasswordCipher, error) {
	privateKey, err := parseRSAPrivateKeyPEM(privateKeyPEM)
	if err != nil {
		return nil, err
	}
	return newPasswordCipherFromKey(privateKey, keyID)
}

func NewEphemeralPasswordCipher(keyID string) (*PasswordCipher, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("generate ephemeral rsa private key: %w", err)
	}
	return newPasswordCipherFromKey(privateKey, keyID)
}

func newPasswordCipherFromKey(privateKey *rsa.PrivateKey, keyID string) (*PasswordCipher, error) {
	if strings.TrimSpace(keyID) == "" {
		keyID = DefaultPasswordCipherKeyID
	}
	publicDER, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, fmt.Errorf("marshal rsa public key: %w", err)
	}
	publicPEM := string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: publicDER}))
	return &PasswordCipher{
		privateKey: privateKey,
		keyID:      keyID,
		publicPEM:  publicPEM,
	}, nil
}

func (c *PasswordCipher) PublicKey() PasswordCipherPublicKey {
	if c == nil {
		return PasswordCipherPublicKey{}
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	return PasswordCipherPublicKey{
		KeyID:     c.keyID,
		Algorithm: "RSA-OAEP-256",
		PublicKey: c.publicPEM,
	}
}

func (c *PasswordCipher) Decrypt(ciphertextB64, keyID string) (string, error) {
	if c == nil {
		return "", fmt.Errorf("password cipher is not configured")
	}
	c.mu.RLock()
	privateKey := c.privateKey
	expectedKeyID := c.keyID
	c.mu.RUnlock()

	if strings.TrimSpace(keyID) != "" && strings.TrimSpace(keyID) != expectedKeyID {
		return "", fmt.Errorf("password cipher key id mismatch")
	}
	ciphertext, err := base64.StdEncoding.DecodeString(strings.TrimSpace(ciphertextB64))
	if err != nil {
		return "", fmt.Errorf("decode encrypted password: %w", err)
	}
	plaintext, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, privateKey, ciphertext, nil)
	if err != nil {
		return "", fmt.Errorf("decrypt encrypted password: %w", err)
	}
	return string(plaintext), nil
}

func parseRSAPrivateKeyPEM(value string) (*rsa.PrivateKey, error) {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return nil, fmt.Errorf("password cipher private key is required")
	}
	block, _ := pem.Decode([]byte(trimmed))
	if block == nil {
		return nil, fmt.Errorf("password cipher private key is not PEM encoded")
	}

	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("parse pkcs1 rsa private key: %w", err)
		}
		return key, nil
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("parse pkcs8 private key: %w", err)
		}
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("private key is not RSA")
		}
		return rsaKey, nil
	default:
		return nil, fmt.Errorf("unsupported private key PEM type %q", block.Type)
	}
}