package encrypt

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io/fs"
	"log"
	"math/big"
	"os"
	"time"
)

// ParsePeerPublicKey umí vzít buď PEM PUBLIC KEY, nebo PEM CERT a vrátí *rsa.PublicKey
func ParsePeerPublicKey(pemOrCert string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemOrCert))
	if block == nil {
		return nil, errf("nenašel jsem PEM blok")
	}
	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, errf("očekávám RSA PUBLIC KEY")
		}
		return rsaPub, nil
	case "CERTIFICATE":
		c, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return nil, err
		}
		rsaPub, ok := c.PublicKey.(*rsa.PublicKey)
		if !ok {
			return nil, errf("cert neobsahuje RSA klíč")
		}
		return rsaPub, nil
	default:
		return nil, errf("nepodporovaný PEM typ: %s", block.Type)
	}
}

func generateSelfSignedCert(pub *rsa.PublicKey, priv *rsa.PrivateKey) []byte {
	tpl := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "Encryptor Local Identity",
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().AddDate(1, 0, 0),
		KeyUsage:  x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		// Tenhle cert je jen „vizitka“ (není pro TLS).
		BasicConstraintsValid: true,
	}
	der, _ := x509.CreateCertificate(rand.Reader, &tpl, &tpl, pub, priv)
	buf := &bytes.Buffer{}
	_ = pem.Encode(buf, &pem.Block{Type: "CERTIFICATE", Bytes: der})
	return buf.Bytes()
}

// malá helper chyba
type strErr string

func (e strErr) Error() string      { return string(e) }
func errf(s string, a ...any) error { return strErr(fmtS(s, a...)) }
func fmtS(format string, a ...any) string {
	var b bytes.Buffer
	b.WriteString(format)
	if len(a) > 0 {
		b.WriteString(": ")
	}
	for i, v := range a {
		if i > 0 {
			b.WriteString(", ")
		}
		b.WriteString(anyToString(v))
	}
	return b.String()
}
func anyToString(v any) string {
	switch t := v.(type) {
	case string:
		return t
	default:
		j, _ := json.Marshal(t)
		return string(j)
	}
}

func LoadOrGenerateKeys(privPath, pubPath, certPath string) (priv *rsa.PrivateKey, pubPEM, certPEM []byte, err error) {

	// existuje private key?
	if fileExists(privPath) && fileExists(pubPath) && fileExists(certPath) {
		// načti privátní klíč
		pkBytes, err := os.ReadFile(privPath)
		if err != nil {
			return nil, nil, nil, err
		}
		block, _ := pem.Decode(pkBytes)
		if block == nil || block.Type != "RSA PRIVATE KEY" {
			return nil, nil, nil, fmt.Errorf("invalid private key PEM")
		}
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, nil, nil, err
		}
		priv = key

		// načti public & cert
		pubPEM, _ = os.ReadFile(pubPath)
		certPEM, _ = os.ReadFile(certPath)

		log.Println("Načtena existující identita z disku.")
		return priv, pubPEM, certPEM, nil
	}

	// jinak vygeneruj nové
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, nil, err
	}
	priv = key

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

	// cert
	certPEM = generateSelfSignedCert(&priv.PublicKey, priv)

	// ulož
	if err := os.WriteFile(privPath,
		pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)}),
		fs.FileMode(0600)); err != nil {
		return nil, nil, nil, err
	}
	if err := os.WriteFile(pubPath, pubPEM, 0644); err != nil {
		return nil, nil, nil, err
	}
	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return nil, nil, nil, err
	}

	log.Println("Vygenerována nová identita a uložena na disk.")
	return priv, pubPEM, certPEM, nil
}

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