package encrypt

// Nový návrh: Password-based šifrované úložiště s derivací klíče (scrypt) +
// vytvoření uživatelských RSA klíčů při prvotní inicializaci. Private/Public
// (a volitelně self-signed cert) se uloží jako položky uvnitř zašifrovaného
// JSON. Soubor není možno otevřít bez hesla.
//
// Formát na disku (JSON):
// {
//   "v":1,
//   "kdf":"scrypt",
//   "N":32768, "r":8, "p":1,
//   "salt":"base64(...)",
//   "n":"base64(nonce)",
//   "c":"base64(ciphertext)"
// }
// Plaintext (po dešifrování):
// {
//   "updated":"RFC3339",
//   "data":{ key: <json>, ... }
// }
// Povinné interní klíče:
//   _identity_private_pem
//   _identity_public_pem
//   _identity_cert_pem   (volitelně)
//
// API: CreateEncryptedStore (nový), OpenEncryptedStore, ChangePassword.
// Původní NewSecureJSONStore (s KeyProvider) je odstraněn – přechod na heslo.

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/scrypt"
)

// SecureJSONStore definuje metody úložiště.
type SecureJSONStore interface {
	Put(key string, value any) error
	Get(key string, dst any) error
	Delete(key string) error
	ListKeys() []string
	Flush() error
	Close() error
	ChangePassword(newPassword string) error
	Has(key string) bool
	// Identity getters (pem). Vrací prázdný string, pokud neexistují.
	IdentityPrivatePEM() string
	IdentityPublicPEM() string
	IdentityCertPEM() string
}

type internalPlain struct {
	Updated time.Time                  `json:"updated"`
	Data    map[string]json.RawMessage `json:"data"`
}

type fileEnvelope struct {
	Version int    `json:"v"`
	KDF     string `json:"kdf"`
	N       int    `json:"N"`
	R       int    `json:"r"`
	P       int    `json:"p"`
	Salt    string `json:"salt"`
	Nonce   string `json:"n"`
	Cipher  string `json:"c"`
}

type secureJSONStore struct {
	mu       sync.RWMutex
	path     string
	kdfN     int
	kdfR     int
	kdfP     int
	salt     []byte
	key      []byte // odvozený AES-256 klíč
	dirty    bool
	plain    internalPlain
	hasIdent bool
}

const (
	defaultKDFN = 32768
	defaultKDFR = 8
	defaultKDFP = 1
)

// CreateEncryptedStore vytvoří nový soubor; error pokud již existuje.
func CreateEncryptedStore(path, password string, generateIdentity bool, commonName string) (SecureJSONStore, error) {
	if err := validatePasswordStrength(password); err != nil {
		return nil, err
	}
	if _, err := os.Stat(path); err == nil {
		return nil, fmt.Errorf("store already exists: %s", path)
	} else if !errors.Is(err, os.ErrNotExist) {
		return nil, err
	}
	s := &secureJSONStore{
		path:  filepath.Clean(path),
		kdfN:  defaultKDFN,
		kdfR:  defaultKDFR,
		kdfP:  defaultKDFP,
		plain: internalPlain{Updated: time.Now().UTC(), Data: make(map[string]json.RawMessage)},
	}
	if err := s.initNew(password, generateIdentity, commonName); err != nil {
		return nil, err
	}
	if err := s.Flush(); err != nil {
		return nil, err
	}
	return s, nil
}

// OpenEncryptedStore načte existující soubor a ověří heslo.
func OpenEncryptedStore(path, password string) (SecureJSONStore, error) {
	if password == "" {
		return nil, errors.New("empty password")
	}
	if _, err := os.Stat(path); err != nil {
		return nil, err
	}
	s := &secureJSONStore{path: filepath.Clean(path)}
	if err := s.load(password); err != nil {
		return nil, err
	}
	return s, nil
}

// initNew vytvoří salt, odvodí klíč a vytvoří identitu pokud je třeba.
func (s *secureJSONStore) initNew(password string, generateIdentity bool, commonName string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return err
	}
	key, err := deriveKey(password, salt, s.kdfN, s.kdfR, s.kdfP)
	if err != nil {
		return err
	}
	s.salt = salt
	s.key = key
	if generateIdentity {
		if err := s.generateIdentityLocked(commonName); err != nil {
			return err
		}
	}
	return nil
}

func (s *secureJSONStore) generateIdentityLocked(commonName string) error {
	// RSA 2048
	pk, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}
	privPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(pk)})
	pubASN1, _ := x509.MarshalPKIXPublicKey(&pk.PublicKey)
	pubPEM := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubASN1})
	// Optional self-signed cert s CommonName uživatele (fallback pokud prázdný)
	if strings.TrimSpace(commonName) == "" {
		commonName = "Local User"
	}
	certPEM, _ := generateSelfSignedCert(pk, commonName)
	_ = s.putLocked("_identity_private_pem", string(privPEM))
	_ = s.putLocked("_identity_public_pem", string(pubPEM))
	_ = s.putLocked("_identity_cert_pem", string(certPEM))
	s.hasIdent = true
	return nil
}

// deriveKey scrypt.
func deriveKey(password string, salt []byte, N, r, p int) ([]byte, error) {
	return scrypt.Key([]byte(password), salt, N, r, p, 32)
}

// load načte soubor, přečte parametry, odvodí klíč a dešifruje.
func (s *secureJSONStore) load(password string) error {
	b, err := os.ReadFile(s.path)
	if err != nil {
		return err
	}
	var env fileEnvelope
	if err := json.Unmarshal(b, &env); err != nil {
		return fmt.Errorf("invalid envelope: %w", err)
	}
	if env.Version != 1 || env.KDF != "scrypt" {
		return errors.New("unsupported version/kdf")
	}
	salt, err := base64.StdEncoding.DecodeString(env.Salt)
	if err != nil {
		return fmt.Errorf("salt b64: %w", err)
	}
	key, err := deriveKey(password, salt, env.N, env.R, env.P)
	if err != nil {
		return err
	}
	nonce, err := base64.StdEncoding.DecodeString(env.Nonce)
	if err != nil {
		return fmt.Errorf("nonce b64: %w", err)
	}
	cBytes, err := base64.StdEncoding.DecodeString(env.Cipher)
	if err != nil {
		return fmt.Errorf("cipher b64: %w", err)
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	plainBytes, err := gcm.Open(nil, nonce, cBytes, nil)
	if err != nil {
		return errors.New("invalid password or corrupted store")
	}
	var pl internalPlain
	if err := json.Unmarshal(plainBytes, &pl); err != nil {
		return fmt.Errorf("plaintext json: %w", err)
	}
	if pl.Data == nil {
		pl.Data = make(map[string]json.RawMessage)
	}
	s.mu.Lock()
	s.kdfN, s.kdfR, s.kdfP = env.N, env.R, env.P
	s.salt = salt
	s.key = key
	s.plain = pl
	s.dirty = false
	s.hasIdent = s.hasIdentityLocked()
	s.mu.Unlock()
	return nil
}

func (s *secureJSONStore) hasIdentityLocked() bool {
	_, ok1 := s.plain.Data["_identity_private_pem"]
	_, ok2 := s.plain.Data["_identity_public_pem"]
	return ok1 && ok2
}

// Put
func (s *secureJSONStore) Put(key string, value any) error {
	if strings.HasPrefix(key, "_identity_") {
		return errors.New("reserved identity key prefix")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.putLocked(key, value)
}

func (s *secureJSONStore) putLocked(key string, value any) error {
	if key == "" {
		return errors.New("empty key")
	}
	b, err := json.Marshal(value)
	if err != nil {
		return err
	}
	if s.plain.Data == nil {
		s.plain.Data = make(map[string]json.RawMessage)
	}
	s.plain.Data[key] = json.RawMessage(b)
	s.plain.Updated = time.Now().UTC()
	s.dirty = true
	return nil
}

// Get
func (s *secureJSONStore) Get(key string, dst any) error {
	s.mu.RLock()
	raw, ok := s.plain.Data[key]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("key not found: %s", key)
	}
	return json.Unmarshal(raw, dst)
}

// Delete
func (s *secureJSONStore) Delete(key string) error {
	if strings.HasPrefix(key, "_identity_") {
		return errors.New("cannot delete identity key")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.plain.Data[key]; !ok {
		return fmt.Errorf("key not found: %s", key)
	}
	delete(s.plain.Data, key)
	s.plain.Updated = time.Now().UTC()
	s.dirty = true
	return nil
}

func (s *secureJSONStore) Has(key string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.plain.Data[key]
	return ok
}

// ListKeys (bez identity interních klíčů)
func (s *secureJSONStore) ListKeys() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make([]string, 0, len(s.plain.Data))
	for k := range s.plain.Data {
		if !strings.HasPrefix(k, "_identity_") {
			out = append(out, k)
		}
	}
	return out
}

// Flush uloží na disk.
func (s *secureJSONStore) Flush() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.dirty {
		return nil
	}
	if len(s.key) == 0 {
		return errors.New("store not initialized")
	}
	plainBytes, err := json.Marshal(s.plain)
	if err != nil {
		return err
	}
	block, err := aes.NewCipher(s.key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return err
	}
	cipherBytes := gcm.Seal(nil, nonce, plainBytes, nil)
	env := fileEnvelope{
		Version: 1,
		KDF:     "scrypt",
		N:       s.kdfN,
		R:       s.kdfR,
		P:       s.kdfP,
		Salt:    base64.StdEncoding.EncodeToString(s.salt),
		Nonce:   base64.StdEncoding.EncodeToString(nonce),
		Cipher:  base64.StdEncoding.EncodeToString(cipherBytes),
	}
	out, err := json.MarshalIndent(env, "", "  ")
	if err != nil {
		return err
	}
	tmp := s.path + ".tmp"
	if err := os.WriteFile(tmp, out, 0o600); err != nil {
		return err
	}
	if err := os.Rename(tmp, s.path); err != nil {
		return err
	}
	s.dirty = false
	return nil
}

func (s *secureJSONStore) Close() error { return s.Flush() }

// ChangePassword re-derivuje klíč a re-encryptne celý soubor atomicky.
func (s *secureJSONStore) ChangePassword(newPassword string) error {
	if err := validatePasswordStrength(newPassword); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	newSalt := make([]byte, 16)
	if _, err := rand.Read(newSalt); err != nil {
		return err
	}
	newKey, err := deriveKey(newPassword, newSalt, s.kdfN, s.kdfR, s.kdfP)
	if err != nil {
		return err
	}
	// Dočasně si zapamatuj starý klíč pro rollback, i když tady rollback neimplementujeme – chyba se projeví před přepsáním.
	oldKey, oldSalt := s.key, s.salt
	s.key, s.salt = newKey, newSalt
	s.dirty = true
	if err := s.Flush(); err != nil {
		// rollback (best effort)
		s.key, s.salt = oldKey, oldSalt
		s.dirty = true
		return err
	}
	return nil
}

// Identity helpers
func (s *secureJSONStore) IdentityPrivatePEM() string { return s.getString("_identity_private_pem") }
func (s *secureJSONStore) IdentityPublicPEM() string  { return s.getString("_identity_public_pem") }
func (s *secureJSONStore) IdentityCertPEM() string    { return s.getString("_identity_cert_pem") }

func (s *secureJSONStore) getString(key string) string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if raw, ok := s.plain.Data[key]; ok {
		var v string
		if err := json.Unmarshal(raw, &v); err == nil {
			return v
		}
	}
	return ""
}

// GenerateRandomPassword vytvoří kryptograficky silné heslo ze sady znaků.
func GenerateRandomPassword(length int) (string, error) {
	if length <= 0 {
		return "", errors.New("invalid length")
	}
	const alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#-_+="
	buf := make([]byte, length)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	for i, b := range buf {
		buf[i] = alphabet[int(b)%len(alphabet)]
	}
	return string(buf), nil
}

// validatePasswordStrength provede základní kontrolu síly hesla.
// Požadavky: min délka 10, alespoň 3 z kategorií (upper, lower, digit, special !@#-_+=)
func validatePasswordStrength(pw string) error {
	if len(pw) < 10 {
		return errors.New("password too short (min 10)")
	}
	var hasU, hasL, hasD, hasS bool
	for _, r := range pw {
		switch {
		case r >= 'A' && r <= 'Z':
			hasU = true
		case r >= 'a' && r <= 'z':
			hasL = true
		case r >= '0' && r <= '9':
			hasD = true
		default:
			if strings.ContainsRune("!@#-_+=", r) {
				hasS = true
			}
		}
	}
	cats := 0
	if hasU {
		cats++
	}
	if hasL {
		cats++
	}
	if hasD {
		cats++
	}
	if hasS {
		cats++
	}
	if cats < 3 {
		return errors.New("password too weak (need 3 types: upper, lower, digit, special)")
	}
	return nil
}

// ValidatePasswordForUI exportuje validaci pro použití ve Fyne UI.
func ValidatePasswordForUI(pw string) error { return validatePasswordStrength(pw) }