package mail

import (
	"crypto/rand"
	"fmt"

	"github.com/ProtonMail/gopenpgp/v2/crypto"
)

type PGPKeyRing struct {
	PrivateKey *crypto.Key
	PublicKey  []byte
}

type PGPService struct {
	keyRing *PGPKeyRing
}

func NewPGPService(privateKeyArmored string) (*PGPService, error) {
	privateKey, err := crypto.NewKeyFromArmored(privateKeyArmored)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	publicKey, err := privateKey.GetPublicKey()
	if err != nil {
		return nil, fmt.Errorf("failed to extract public key: %w", err)
	}

	return &PGPService{
		keyRing: &PGPKeyRing{
			PrivateKey: privateKey,
			PublicKey:  publicKey,
		},
	}, nil
}

func (s *PGPService) Encrypt(plaintext string, recipientPublicKey *crypto.Key) (string, error) {
	pgpMessage, err := crypto.NewPlainMessage([]byte(plaintext))
	if err != nil {
		return "", fmt.Errorf("failed to create PGP message: %w", err)
	}

	encrypted, err := pgpMessage.Encrypt(recipientPublicKey)
	if err != nil {
		return "", fmt.Errorf("failed to encrypt: %w", err)
	}

	return encrypted.GetArmored()
}

func (s *PGPService) EncryptAndSign(plaintext string, recipientPublicKey *crypto.Key, passphrase string) (string, error) {
	pgpMessage, err := crypto.NewPlainMessage([]byte(plaintext))
	if err != nil {
		return "", fmt.Errorf("failed to create PGP message: %w", err)
	}

	encrypted, err := pgpMessage.EncryptAndSign(recipientPublicKey, s.keyRing.PrivateKey, []byte(passphrase))
	if err != nil {
		return "", fmt.Errorf("failed to encrypt and sign: %w", err)
	}

	return encrypted.GetArmored()
}

func (s *PGPService) Decrypt(encrypted string, passphrase string) (string, error) {
	armoredKey, err := crypto.NewKeyFromArmored(encrypted)
	if err != nil {
		return "", fmt.Errorf("failed to parse armored key: %w", err)
	}

	pgpMessage, err := crypto.NewPlainMessageFromString(armoredKey.GetArmored())
	if err != nil {
		return "", fmt.Errorf("failed to parse encrypted message: %w", err)
	}

	decrypted, err := pgpMessage.Decrypt(s.keyRing.PrivateKey, []byte(passphrase))
	if err != nil {
		return "", fmt.Errorf("failed to decrypt: %w", err)
	}

	return string(decrypted.GetBinary()), nil
}

func (s *PGPService) GenerateKeyPair(email string, passphrase string) (privateKey, publicKey string, err error) {
	key, err := crypto.GenerateKey(email, passphrase, "RSA", 4096)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate key pair: %w", err)
	}

	privateArmor, err := key.Armor()
	if err != nil {
		return "", "", fmt.Errorf("failed to armor private key: %w", err)
	}

	pubKeyBytes, err := key.GetPublicKey()
	if err != nil {
		return "", "", fmt.Errorf("failed to extract public key: %w", err)
	}

	pubArmor := string(pubKeyBytes)

	return string(privateArmor), pubArmor, nil
}

func (s *PGPService) GetFingerprint() (string, error) {
	if s.keyRing == nil || s.keyRing.PrivateKey == nil {
		return "", fmt.Errorf("no key ring available")
	}
	fingerprint := s.keyRing.PrivateKey.GetFingerprint()
	return fingerprint, nil
}

func (s *PGPService) SignData(data []byte, passphrase string) (string, error) {
	pgpMessage, err := crypto.NewPlainMessage(data)
	if err != nil {
		return "", fmt.Errorf("failed to create PGP message: %w", err)
	}

	signed, err := pgpMessage.Sign(s.keyRing.PrivateKey, []byte(passphrase))
	if err != nil {
		return "", fmt.Errorf("failed to sign data: %w", err)
	}

	return signed.GetArmored()
}

func (s *PGPService) EncryptAttachment(data []byte, recipientPublicKey *crypto.Key) (*Attachment, error) {
	symKey := make([]byte, 32)
	if _, err := rand.Read(symKey); err != nil {
		return nil, fmt.Errorf("failed to generate symmetric key: %w", err)
	}

	symKeyRing, err := crypto.NewKeyFromArmored(recipientPublicKey.GetArmored())
	if err != nil {
		return nil, fmt.Errorf("failed to parse recipient key: %w", err)
	}

	pgpMessage, err := crypto.NewPlainMessage(data)
	if err != nil {
		return nil, fmt.Errorf("failed to create PGP message: %w", err)
	}

	encrypted, err := pgpMessage.Encrypt(symKeyRing)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt attachment: %w", err)
	}

	encData := []byte(encrypted.GetBinary())

	encryptedSymKey, err := symKeyRing.Encrypt(symKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt symmetric key: %w", err)
	}

	return &Attachment{
		DataEnc: string(encData),
		Keys: []AttachmentKey{{
			DataEnc: string(encryptedSymKey.GetBinary()),
		}},
	}, nil
}

func (s *PGPService) DecryptAttachment(attachment *Attachment, passphrase string) ([]byte, error) {
	if len(attachment.Keys) == 0 {
		return nil, fmt.Errorf("no keys available for attachment decryption")
	}

	encryptedSymKey, err := crypto.NewKeyFromArmored(string(attachment.Keys[0].DataEnc))
	if err != nil {
		return nil, fmt.Errorf("failed to parse encrypted symmetric key: %w", err)
	}

	symKey, err := encryptedSymKey.Decrypt([]byte(passphrase))
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt symmetric key: %w", err)
	}

	pgpMessage, err := crypto.NewPlainMessage([]byte(attachment.DataEnc))
	if err != nil {
		return nil, fmt.Errorf("failed to create PGP message: %w", err)
	}

	decrypted, err := pgpMessage.DecryptWithKey(s.keyRing.PrivateKey, symKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt attachment: %w", err)
	}

	return decrypted.GetBinary(), nil
}