package encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io/fs"
	"math/big"
	"os"
	"path/filepath"
	"time"
)

type Envelope struct {
	EK string `json:"ek"`
	N  string `json:"n"`
	CT string `json:"ct"`
}

type Service struct {
	priv    *rsa.PrivateKey
	pubPEM  []byte
	certPEM []byte
	dir     string
}

func NewService(storageDir string) (*Service, error) {
	s := &Service{dir: storageDir}
	if err := s.loadOrGenerateKeys(); err != nil {
		return nil, err
	}
	return s, nil
}

func (s *Service) PublicPEM() string  { return string(s.pubPEM) }
func (s *Service) PublicCert() string { return string(s.certPEM) }

func (s *Service) Encrypt(message, peerPEMorCert string) (string, error) {
	pubKey, err := parsePeerPublicKey(peerPEMorCert)
	if err != nil {
		return "", err
	}

	aesKey := make([]byte, 32)
	if _, err := rand.Read(aesKey); err != nil {
		return "", err
	}

	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}

	ciphertext := gcm.Seal(nil, nonce, []byte(message), nil)
	label := []byte{}
	ek, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, aesKey, label)
	if err != nil {
		return "", err
	}

	env := Envelope{
		EK: base64.StdEncoding.EncodeToString(ek),
		N:  base64.StdEncoding.EncodeToString(nonce),
		CT: base64.StdEncoding.EncodeToString(ciphertext),
	}
	out, _ := json.MarshalIndent(env, "", "  ")
	return string(out), nil
}

func (s *Service) Decrypt(payload string) (string, error) {
	var env Envelope
	if err := json.Unmarshal([]byte(payload), &env); err != nil {
		return "", fmt.Errorf("invalid JSON: %w", err)
	}
	ek, err := base64.StdEncoding.DecodeString(env.EK)
	if err != nil {
		return "", fmt.Errorf("ek b64: %w", err)
	}
	nonce, err := base64.StdEncoding.DecodeString(env.N)
	if err != nil {
		return "", fmt.Errorf("n b64: %w", err)
	}
	ct, err := base64.StdEncoding.DecodeString(env.CT)
	if err != nil {
		return "", fmt.Errorf("ct b64: %w", err)
	}

	label := []byte{}
	aesKey, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, s.priv, ek, label)
	if err != nil {
		return "", fmt.Errorf("RSA-OAEP decrypt: %w", err)
	}

	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	plain, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		return "", fmt.Errorf("GCM open: %w", err)
	}
	return string(plain), nil
}

func parsePeerPublicKey(pemOrCert string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemOrCert))
	if block == nil {
		return nil, errors.New("no PEM block found")
	}
	switch block.Type {
	case "PUBLIC KEY":
		k, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		rsaPub, ok := k.(*rsa.PublicKey)
		if !ok {
			return nil, errors.New("expecting RSA PUBLIC KEY")
		}
		return rsaPub, nil
	case "RSA PUBLIC KEY":
		return x509.ParsePKCS1PublicKey(block.Bytes)
	case "CERTIFICATE":
		c, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return nil, err
		}
		rsaPub, ok := c.PublicKey.(*rsa.PublicKey)
		if !ok {
			return nil, errors.New("certificate does not contain RSA key")
		}
		return rsaPub, nil
	default:
		return nil, fmt.Errorf("unsupported PEM type: %s", block.Type)
	}
}

func ParsePeerPublicKey(pemOrCert string) (*rsa.PublicKey, error) {
	return parsePeerPublicKey(pemOrCert)
}

func (s *Service) loadOrGenerateKeys() error {
	privPath := filepath.Join(s.dir, "identity_key.pem")
	pubPath := filepath.Join(s.dir, "public.pem")
	certPath := filepath.Join(s.dir, "identity.crt")

	if fileExists(privPath) && fileExists(pubPath) && fileExists(certPath) {
		pkBytes, err := os.ReadFile(privPath)
		if err != nil {
			return err
		}
		block, _ := pem.Decode(pkBytes)
		if block == nil || block.Type != "RSA PRIVATE KEY" {
			return errors.New("invalid private key PEM")
		}
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return err
		}
		s.priv = key
		s.pubPEM, _ = os.ReadFile(pubPath)
		s.certPEM, _ = os.ReadFile(certPath)
		return nil
	}

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}
	s.priv = key

	pubASN1, _ := x509.MarshalPKIXPublicKey(&s.priv.PublicKey)
	s.pubPEM = pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubASN1})

	cert, err := generateSelfSignedCert(s.priv)
	if err != nil {
		return err
	}
	s.certPEM = cert

	if err := os.MkdirAll(s.dir, 0o700); err != nil {
		return err
	}
	if err := os.WriteFile(privPath, pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(s.priv)}), fs.FileMode(0o600)); err != nil {
		return err
	}
	if err := os.WriteFile(pubPath, s.pubPEM, 0o644); err != nil {
		return err
	}
	if err := os.WriteFile(certPath, s.certPEM, 0o644); err != nil {
		return err
	}
	return nil
}

func generateSelfSignedCert(priv *rsa.PrivateKey) ([]byte, error) {
	tpl := &x509.Certificate{
		SerialNumber:          big.NewInt(time.Now().UnixNano()),
		Subject:               pkix.Name{CommonName: "Encryptor Local Identity"},
		NotBefore:             time.Now().Add(-1 * time.Hour),
		NotAfter:              time.Now().AddDate(1, 0, 0),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
	}
	der, err := x509.CreateCertificate(rand.Reader, tpl, tpl, &priv.PublicKey, priv)
	if err != nil {
		return nil, err
	}
	buf := &bytes.Buffer{}
	_ = pem.Encode(buf, &pem.Block{Type: "CERTIFICATE", Bytes: der})
	return buf.Bytes(), nil
}

func fileExists(p string) bool {
	info, err := os.Stat(p)
	return err == nil && !info.IsDir()
}