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 := normalizePrivateKeyPEM(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)
	}
}

func normalizePrivateKeyPEM(value string) string {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return ""
	}
	normalized := strings.ReplaceAll(trimmed, "\r\n", "\n")
	normalized = strings.ReplaceAll(normalized, "\r", "\n")
	normalized = strings.ReplaceAll(normalized, `\n`, "\n")
	if block, _ := pem.Decode([]byte(normalized)); block != nil {
		return normalized
	}

	for _, blockType := range []string{"RSA PRIVATE KEY", "PRIVATE KEY"} {
		begin := "-----BEGIN " + blockType + "-----"
		end := "-----END " + blockType + "-----"
		start := strings.Index(normalized, begin)
		stop := strings.Index(normalized, end)
		if start < 0 || stop <= start {
			continue
		}
		body := normalized[start+len(begin) : stop]
		body = strings.Join(strings.Fields(body), "")
		if body == "" {
			return normalized
		}
		return begin + "\n" + wrapBase64Lines(body, 64) + "\n" + end + "\n"
	}
	return normalized
}

func wrapBase64Lines(value string, width int) string {
	if width <= 0 || len(value) <= width {
		return value
	}
	var b strings.Builder
	b.Grow(len(value) + len(value)/width)
	for len(value) > width {
		b.WriteString(value[:width])
		b.WriteByte('\n')
		value = value[width:]
	}
	b.WriteString(value)
	return b.String()
}