package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	encrypt "fckeuspy-go/lib"
	"fmt"
	mrand "math/rand"
)

// VaultService implementuje ServiceFacade nad SecureJSONStore.
// Využívá existující hybridní šifrování z encrypt.Service (re-use kódu voláním helperů?)
// Pro jednoduchost zopakujeme potřebnou část: Encrypt/Decrypt využijí dočasný encrypt.Service
// inicializovaný z načtených klíčů.
type VaultService struct {
	store   encrypt.SecureJSONStore
	priv    *rsa.PrivateKey
	pubPEM  string
	certPEM string
}

func NewVaultService(store encrypt.SecureJSONStore) (*VaultService, error) {
	vs := &VaultService{store: store}
	if err := vs.loadIdentity(); err != nil {
		return nil, err
	}
	return vs, nil
}

func (v *VaultService) loadIdentity() error {
	privPem := v.store.IdentityPrivatePEM()
	if privPem == "" {
		return errors.New("missing private key in store")
	}
	block, _ := pem.Decode([]byte(privPem))
	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
	}
	v.priv = key
	v.pubPEM = v.store.IdentityPublicPEM()
	v.certPEM = v.store.IdentityCertPEM()
	if v.pubPEM == "" { // derive public if missing
		pubASN1, _ := x509.MarshalPKIXPublicKey(&key.PublicKey)
		v.pubPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubASN1}))
	}
	return nil
}

func (v *VaultService) PublicPEM() string  { return v.pubPEM }
func (v *VaultService) PublicCert() string { return v.certPEM }

// Encrypt provede hybridní šifrování stejně jako původní Service.Encrypt.
func (v *VaultService) Encrypt(message, peerPEMorCert string) (string, error) {
	// Vytvoříme ad-hoc Service s privátním klíčem jen pro volání existující logiky? Nejjednodušší je zkopírovat potřebnou část kódu.
	return encryptHybrid(v.priv, message, peerPEMorCert)
}

// Decrypt provede rozšifrování.
func (v *VaultService) Decrypt(payload string) (string, error) { return decryptHybrid(v.priv, payload) }

// --- Lokální helpery (duplikace z encrypt.Service, zredukované) ---

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

// --- Contacts management ---

type Contact struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	Cert string `json:"cert"` // full PEM cert (or public key PEM)
}

const contactsKey = "contacts"

func (v *VaultService) ListContacts() ([]Contact, error) {
	var list []Contact
	if !v.store.Has(contactsKey) {
		return []Contact{}, nil
	}
	if err := v.store.Get(contactsKey, &list); err != nil {
		return nil, err
	}
	return list, nil
}

func (v *VaultService) SaveContact(c Contact) error {
	list, _ := v.ListContacts()
	// upsert by ID; if empty ID, assign a new unique random ID to avoid overwriting
	if c.ID == "" {
		c.ID = generateUniqueContactID(list)
	}
	replaced := false
	for i := range list {
		if list[i].ID == c.ID {
			list[i] = c
			replaced = true
			break
		}
	}
	if !replaced {
		list = append(list, c)
	}
	if err := v.store.Put(contactsKey, list); err != nil {
		return err
	}
	return v.store.Flush()
}

func (v *VaultService) DeleteContact(id string) error {
	list, _ := v.ListContacts()
	out := list[:0]
	for _, c := range list {
		if c.ID != id {
			out = append(out, c)
		}
	}
	if err := v.store.Put(contactsKey, out); err != nil {
		return err
	}
	return v.store.Flush()
}

func deriveContactID(c Contact) string {
	// simple stable ID: prefer CN from cert; fallback to sha256 of cert
	if cn := extractCN(c.Cert); cn != "" {
		return cn
	}
	sum := sha256.Sum256([]byte(c.Cert))
	return fmt.Sprintf("%x", sum[:8])
}

// generateUniqueContactID creates a short random hex ID that does not collide with existing list.
func generateUniqueContactID(existing []Contact) string {
	exists := make(map[string]struct{}, len(existing))
	for _, c := range existing {
		exists[c.ID] = struct{}{}
	}
	for {
		b := make([]byte, 6)
		if _, err := rand.Read(b); err != nil {
			// fallback to time-based seed rand
			n := mrand.Int63()
			return fmt.Sprintf("%x", n)
		}
		id := fmt.Sprintf("%x", b)
		if _, ok := exists[id]; !ok {
			return id
		}
	}
}

func extractCN(pemText string) string {
	// Support multiple concatenated PEM blocks (public key + cert, or chain)
	rest := []byte(pemText)
	for len(rest) > 0 {
		var block *pem.Block
		block, rest = pem.Decode(rest)
		if block == nil {
			break
		}
		if block.Type == "CERTIFICATE" {
			if cert, err := x509.ParseCertificate(block.Bytes); err == nil {
				return cert.Subject.CommonName
			}
		}
		// continue to next block until certificate found
	}
	return ""
}

func encryptHybrid(priv *rsa.PrivateKey, message, peerPEMorCert string) (string, error) {
	pubKey, err := encrypt.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
	}
	ct := gcm.Seal(nil, nonce, []byte(message), nil)
	ek, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, aesKey, []byte{})
	if err != nil {
		return "", err
	}
	env := hybridEnvelope{EK: base64.StdEncoding.EncodeToString(ek), N: base64.StdEncoding.EncodeToString(nonce), CT: base64.StdEncoding.EncodeToString(ct)}
	out, _ := json.MarshalIndent(env, "", "  ")
	return string(out), nil
}

func decryptHybrid(priv *rsa.PrivateKey, payload string) (string, error) {
	var env hybridEnvelope
	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)
	}
	aesKey, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, priv, ek, []byte{})
	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
	}
	pt, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		return "", fmt.Errorf("GCM open: %w", err)
	}
	return string(pt), nil
}