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 `json:"priv,omitempty"`
	PubPEM  []byte          `json:"pub,omitempty"`
	CertPEM []byte          `json:"cert,omitempty"`
	dir     string          `json:"-"`
	symKey  []byte          `json:"-"` // interní AES klíč pro lokální šifrování
}

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")
	symPath := filepath.Join(s.dir, "sym.key") // base64(klic)

	fmt.Printf("Using storage dir: %s\n", s.dir)
	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)
		// Načti symetrický klíč (pokud existuje); pokud ne, vytvoř.
		if fileExists(symPath) {
			b, err := os.ReadFile(symPath)
			if err != nil {
				return err
			}
			decoded, err := base64.StdEncoding.DecodeString(string(b))
			if err != nil {
				return fmt.Errorf("sym.key base64 decode: %w", err)
			}
			if l := len(decoded); l != 16 && l != 24 && l != 32 {
				return fmt.Errorf("unsupported sym key length: %d", l)
			}
			s.symKey = decoded
		} else {
			if err := s.generateAndPersistSymKey(symPath); err != nil {
				return err
			}
		}
		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
	}
	// Vytvoř symetrický klíč
	if err := s.generateAndPersistSymKey(symPath); err != nil {
		return err
	}
	return nil
}

// generateAndPersistSymKey vytvoří nový AES-256 klíč a uloží jej (pokud existuje dir).
func (s *Service) generateAndPersistSymKey(path string) error {
	k := make([]byte, 32)
	if _, err := rand.Read(k); err != nil {
		return err
	}
	s.symKey = k
	if s.dir == "" { // ephemeral (např. web mód bez persistence)
		return nil
	}
	enc := base64.StdEncoding.EncodeToString(k)
	return os.WriteFile(path, []byte(enc), 0o600)
}

// SymEnvelope je malý formát pro interní symetrické šifrování.
type SymEnvelope struct {
	N  string `json:"n"`  // base64(nonce)
	CT string `json:"ct"` // base64(ciphertext||tag)
}

// SymmetricEncrypt šifruje plaintext pomocí interního AES-GCM klíče.
func (s *Service) SymmetricEncrypt(plain string) (string, error) {
	if len(s.symKey) == 0 {
		return "", errors.New("sym key not initialized")
	}
	block, err := aes.NewCipher(s.symKey)
	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
	}
	ct := gcm.Seal(nil, nonce, []byte(plain), nil)
	env := SymEnvelope{
		N:  base64.StdEncoding.EncodeToString(nonce),
		CT: base64.StdEncoding.EncodeToString(ct),
	}
	out, _ := json.Marshal(env)
	return string(out), nil
}

// SymmetricDecrypt dešifruje JSON envelope vytvořený SymmetricEncrypt.
func (s *Service) SymmetricDecrypt(payload string) (string, error) {
	if len(s.symKey) == 0 {
		return "", errors.New("sym key not initialized")
	}
	var env SymEnvelope
	if err := json.Unmarshal([]byte(payload), &env); err != nil {
		return "", fmt.Errorf("invalid sym envelope json: %w", err)
	}
	nonce, err := base64.StdEncoding.DecodeString(env.N)
	if err != nil {
		return "", fmt.Errorf("nonce b64: %w", err)
	}
	ct, err := base64.StdEncoding.DecodeString(env.CT)
	if err != nil {
		return "", fmt.Errorf("ct b64: %w", err)
	}
	block, err := aes.NewCipher(s.symKey)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	pt, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		return "", fmt.Errorf("gcm open: %w", err)
	}
	return string(pt), 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()
}