package keymanagement

import (
    "crypto/rand"
    "encoding/hex"
    "fmt"
    "io"
)

// GenerateRandomBytes generates cryptographically secure random bytes
func GenerateRandomBytes(length int) ([]byte, error) {
    bytes := make([]byte, length)
    if _, err := io.ReadFull(rand.Reader, bytes); err != nil {
        return nil, fmt.Errorf("failed to generate random bytes: %w", err)
    }
    return bytes, nil
}

// GenerateEntropy generates entropy for key generation
func GenerateEntropy(bits int) ([]byte, error) {
    if bits%8 != 0 {
        return nil, fmt.Errorf("bits must be divisible by 8")
    }

    entropy, err := GenerateRandomBytes(bits / 8)
    if err != nil {
        return nil, err
    }

    return entropy, nil
}

// Multi-algorithm key generation
type KeyType int

const (
    KeyTypeSecp256k1 KeyType = iota
    KeyTypeP256
    KeyTypeBLS
    KeyTypeMLDSA
    KeyTypeMLKEM
    KeyTypeSLHDSA
)

type KeyPair struct {
    Type       KeyType
    PrivateKey interface{}
    PublicKey  interface{}
}

func GenerateKeyPair(keyType KeyType) (*KeyPair, error) {
    switch keyType {
    case KeyTypeSecp256k1:
        key, err := ecdsa.GenerateKey(secp256k1.S256(), rand.Reader)
        if err != nil {
            return nil, err
        }
        return &KeyPair{
            Type:       KeyTypeSecp256k1,
            PrivateKey: key,
            PublicKey:  &key.PublicKey,
        }, nil

    case KeyTypeBLS:
        privKey, err := bls.NewSecretKey()
        if err != nil {
            return nil, err
        }
        return &KeyPair{
            Type:       KeyTypeBLS,
            PrivateKey: privKey,
            PublicKey:  bls.PublicKeyFromSecretKey(privKey),
        }, nil

    case KeyTypeMLDSA:
        pubKey, privKey, err := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
        if err != nil {
            return nil, err
        }
        return &KeyPair{
            Type:       KeyTypeMLDSA,
            PrivateKey: privKey,
            PublicKey:  pubKey,
        }, nil

    default:
        return nil, fmt.Errorf("unsupported key type: %v", keyType)
    }
}

import (
    "github.com/tyler-smith/go-bip32"
    "github.com/tyler-smith/go-bip39"
)

// HDWallet manages hierarchical deterministic keys
type HDWallet struct {
    seed       []byte
    masterKey  *bip32.Key
    mnemonic   string
    passphrase string
}

// NewHDWallet creates a new HD wallet from mnemonic
func NewHDWallet(mnemonic, passphrase string) (*HDWallet, error) {
    // Validate mnemonic
    if !bip39.IsMnemonicValid(mnemonic) {
        return nil, fmt.Errorf("invalid mnemonic")
    }

    // Generate seed
    seed := bip39.NewSeed(mnemonic, passphrase)

    // Generate master key
    masterKey, err := bip32.NewMasterKey(seed)
    if err != nil {
        return nil, err
    }

    return &HDWallet{
        seed:       seed,
        masterKey:  masterKey,
        mnemonic:   mnemonic,
        passphrase: passphrase,
    }, nil
}

// GenerateNewWallet creates a new wallet with random mnemonic
func GenerateNewWallet(words int, passphrase string) (*HDWallet, error) {
    // Generate entropy
    var entropy []byte
    var err error

    switch words {
    case 12:
        entropy, err = GenerateEntropy(128)
    case 15:
        entropy, err = GenerateEntropy(160)
    case 18:
        entropy, err = GenerateEntropy(192)
    case 21:
        entropy, err = GenerateEntropy(224)
    case 24:
        entropy, err = GenerateEntropy(256)
    default:
        return nil, fmt.Errorf("invalid word count: must be 12, 15, 18, 21, or 24")
    }

    if err != nil {
        return nil, err
    }

    // Generate mnemonic
    mnemonic, err := bip39.NewMnemonic(entropy)
    if err != nil {
        return nil, err
    }

    return NewHDWallet(mnemonic, passphrase)
}

// DeriveAccount derives an account key using BIP-44 path
func (w *HDWallet) DeriveAccount(coinType, account uint32) (*bip32.Key, error) {
    // m/44'/coinType'/account'
    path := []uint32{
        bip32.FirstHardenedChild + 44,       // purpose
        bip32.FirstHardenedChild + coinType, // coin type
        bip32.FirstHardenedChild + account,  // account
    }

    key := w.masterKey
    for _, index := range path {
        var err error
        key, err = key.NewChildKey(index)
        if err != nil {
            return nil, err
        }
    }

    return key, nil
}

// DeriveAddress derives a specific address
func (w *HDWallet) DeriveAddress(coinType, account, change, index uint32) (*ecdsa.PrivateKey, string, error) {
    // Get account key
    accountKey, err := w.DeriveAccount(coinType, account)
    if err != nil {
        return nil, "", err
    }

    // Derive change key
    changeKey, err := accountKey.NewChildKey(change)
    if err != nil {
        return nil, "", err
    }

    // Derive address key
    addressKey, err := changeKey.NewChildKey(index)
    if err != nil {
        return nil, "", err
    }

    // Convert to ECDSA key
    privateKey := &ecdsa.PrivateKey{
        PublicKey: ecdsa.PublicKey{
            Curve: secp256k1.S256(),
        },
        D: new(big.Int).SetBytes(addressKey.Key),
    }

    privateKey.PublicKey.X, privateKey.PublicKey.Y = privateKey.Curve.ScalarBaseMult(addressKey.Key)

    // Generate address based on coin type
    var address string
    switch coinType {
    case 0: // Bitcoin
        address = bitcoinAddress(&privateKey.PublicKey)
    case 60: // Ethereum
        address = ethereumAddress(&privateKey.PublicKey)
    default:
        address = hex.EncodeToString(addressKey.PublicKey())
    }

    return privateKey, address, nil
}

import (
    "golang.org/x/crypto/argon2"
    "crypto/chacha20poly1305"
)

// KeyStore manages encrypted key storage
type KeyStore struct {
    keys map[string]*EncryptedKey
    salt []byte
}

type EncryptedKey struct {
    ID         string
    Type       KeyType
    Ciphertext []byte
    Nonce      []byte
    Salt       []byte
    Params     *KDFParams
}

type KDFParams struct {
    Algorithm string
    Time      uint32
    Memory    uint32
    Threads   uint8
    KeyLen    uint32
}

// EncryptKey encrypts a private key with a password
func EncryptKey(key []byte, password string) (*EncryptedKey, error) {
    // Generate salt
    salt, err := GenerateRandomBytes(32)
    if err != nil {
        return nil, err
    }

    // Derive encryption key using Argon2id
    params := &KDFParams{
        Algorithm: "argon2id",
        Time:      1,
        Memory:    64 * 1024,
        Threads:   4,
        KeyLen:    32,
    }

    encKey := argon2.IDKey(
        []byte(password),
        salt,
        params.Time,
        params.Memory,
        params.Threads,
        params.KeyLen,
    )

    // Encrypt with ChaCha20-Poly1305
    aead, err := chacha20poly1305.New(encKey)
    if err != nil {
        return nil, err
    }

    nonce, err := GenerateRandomBytes(aead.NonceSize())
    if err != nil {
        return nil, err
    }

    ciphertext := aead.Seal(nil, nonce, key, nil)

    return &EncryptedKey{
        ID:         hex.EncodeToString(GenerateRandomBytes(16)),
        Ciphertext: ciphertext,
        Nonce:      nonce,
        Salt:       salt,
        Params:     params,
    }, nil
}

// DecryptKey decrypts an encrypted private key
func DecryptKey(encKey *EncryptedKey, password string) ([]byte, error) {
    // Derive decryption key
    decKey := argon2.IDKey(
        []byte(password),
        encKey.Salt,
        encKey.Params.Time,
        encKey.Params.Memory,
        encKey.Params.Threads,
        encKey.Params.KeyLen,
    )

    // Decrypt with ChaCha20-Poly1305
    aead, err := chacha20poly1305.New(decKey)
    if err != nil {
        return nil, err
    }

    plaintext, err := aead.Open(nil, encKey.Nonce, encKey.Ciphertext, nil)
    if err != nil {
        return nil, fmt.Errorf("decryption failed: %w", err)
    }

    return plaintext, nil
}

// HSMKeyManager interfaces with hardware security modules
type HSMKeyManager struct {
    session HSMSession
    pins    map[string]string
}

// GenerateKeyInHSM generates a key within the HSM
func (h *HSMKeyManager) GenerateKeyInHSM(keyType KeyType, label string) (string, error) {
    var mechanism uint
    var keyTemplate []Attribute

    switch keyType {
    case KeyTypeSecp256k1:
        mechanism = CKM_EC_KEY_PAIR_GEN
        keyTemplate = []Attribute{
            {Type: CKA_EC_PARAMS, Value: secp256k1Params},
            {Type: CKA_TOKEN, Value: true},
            {Type: CKA_PRIVATE, Value: true},
            {Type: CKA_SENSITIVE, Value: true},
            {Type: CKA_EXTRACTABLE, Value: false},
            {Type: CKA_LABEL, Value: label},
        }

    case KeyTypeBLS:
        // Custom BLS implementation for HSM
        mechanism = CKM_VENDOR_BLS_KEY_GEN
        keyTemplate = []Attribute{
            {Type: CKA_KEY_TYPE, Value: CKK_BLS},
            {Type: CKA_TOKEN, Value: true},
            {Type: CKA_SENSITIVE, Value: true},
            {Type: CKA_LABEL, Value: label},
        }

    default:
        return "", fmt.Errorf("unsupported key type for HSM")
    }

    keyHandle, err := h.session.GenerateKeyPair(mechanism, keyTemplate)
    if err != nil {
        return "", err
    }

    return keyHandle, nil
}

// SignWithHSM signs data using a key stored in HSM
func (h *HSMKeyManager) SignWithHSM(keyHandle string, data []byte) ([]byte, error) {
    return h.session.Sign(keyHandle, data)
}

import (
    "golang.org/x/crypto/hkdf"
    "golang.org/x/crypto/pbkdf2"
    "golang.org/x/crypto/scrypt"
)

// KDF interface for different key derivation functions
type KDF interface {
    DeriveKey(password []byte, salt []byte, keyLen int) ([]byte, error)
}

// PBKDF2 implementation
type PBKDF2 struct {
    Iterations int
    HashFunc   func() hash.Hash
}

func (p *PBKDF2) DeriveKey(password, salt []byte, keyLen int) ([]byte, error) {
    return pbkdf2.Key(password, salt, p.Iterations, keyLen, p.HashFunc), nil
}

// Scrypt implementation
type Scrypt struct {
    N      int // CPU/memory cost
    R      int // Block size
    P      int // Parallelization
    KeyLen int
}

func (s *Scrypt) DeriveKey(password, salt []byte, keyLen int) ([]byte, error) {
    return scrypt.Key(password, salt, s.N, s.R, s.P, keyLen)
}

// Argon2 implementation
type Argon2 struct {
    Time    uint32
    Memory  uint32
    Threads uint8
    Type    string // "argon2i", "argon2d", or "argon2id"
}

func (a *Argon2) DeriveKey(password, salt []byte, keyLen int) ([]byte, error) {
    switch a.Type {
    case "argon2i":
        return argon2.Key(password, salt, a.Time, a.Memory, a.Threads, uint32(keyLen)), nil
    case "argon2id":
        return argon2.IDKey(password, salt, a.Time, a.Memory, a.Threads, uint32(keyLen)), nil
    default:
        return nil, fmt.Errorf("unsupported argon2 type: %s", a.Type)
    }
}

// HKDF for key expansion
func ExpandKey(secret, salt, info []byte, length int) ([]byte, error) {
    hkdf := hkdf.New(sha256.New, secret, salt, info)
    key := make([]byte, length)
    if _, err := io.ReadFull(hkdf, key); err != nil {
        return nil, err
    }
    return key, nil
}

import (
    "github.com/hashicorp/vault/shamir"
)

// SecretSharing implements Shamir's secret sharing
type SecretSharing struct {
    threshold int
    shares    int
}

// SplitSecret splits a secret into n shares with k threshold
func (ss *SecretSharing) SplitSecret(secret []byte) ([][]byte, error) {
    return shamir.Split(secret, ss.shares, ss.threshold)
}

// RecoverSecret recovers the secret from shares
func (ss *SecretSharing) RecoverSecret(shares [][]byte) ([]byte, error) {
    return shamir.Combine(shares)
}

// Multi-party key generation
type MultiPartyKeyGen struct {
    participants int
    threshold    int
    shares       map[string][][]byte
}

func (m *MultiPartyKeyGen) GenerateDistributedKey() error {
    // Generate master key
    masterKey, err := GenerateRandomBytes(32)
    if err != nil {
        return err
    }

    // Split using Shamir's secret sharing
    shares, err := shamir.Split(masterKey, m.participants, m.threshold)
    if err != nil {
        return err
    }

    // Distribute shares to participants
    for i, share := range shares {
        participantID := fmt.Sprintf("participant_%d", i)
        m.shares[participantID] = append(m.shares[participantID], share)
    }

    // Clear master key from memory
    for i := range masterKey {
        masterKey[i] = 0
    }

    return nil
}

import (
    "time"
)

// KeyRotationPolicy defines rotation rules
type KeyRotationPolicy struct {
    MaxAge          time.Duration
    MaxOperations   int64
    ForceRotateTime time.Time
}

// KeyRotationManager handles automatic key rotation
type KeyRotationManager struct {
    policy      *KeyRotationPolicy
    currentKey  *KeyPair
    keyHistory  []*KeyPair
    operations  int64
    createdAt   time.Time
}

func (krm *KeyRotationManager) ShouldRotate() bool {
    // Check age
    if time.Since(krm.createdAt) > krm.policy.MaxAge {
        return true
    }

    // Check operations count
    if krm.operations >= krm.policy.MaxOperations {
        return true
    }

    // Check force rotation time
    if time.Now().After(krm.policy.ForceRotateTime) {
        return true
    }

    return false
}

func (krm *KeyRotationManager) RotateKey() (*KeyPair, error) {
    // Generate new key
    newKey, err := GenerateKeyPair(krm.currentKey.Type)
    if err != nil {
        return nil, err
    }

    // Archive current key
    krm.keyHistory = append(krm.keyHistory, krm.currentKey)

    // Update current key
    krm.currentKey = newKey
    krm.operations = 0
    krm.createdAt = time.Now()

    // Trigger re-encryption of data with new key
    go krm.reencryptData(krm.keyHistory[len(krm.keyHistory)-1], newKey)

    return newKey, nil
}

func (krm *KeyRotationManager) reencryptData(oldKey, newKey *KeyPair) {
    // Re-encrypt all data encrypted with old key
    // This is application-specific
}

// MnemonicRecovery handles key recovery from mnemonic phrases
type MnemonicRecovery struct {
    wordList []string
}

// RecoverFromMnemonic recovers keys from a mnemonic phrase
func (mr *MnemonicRecovery) RecoverFromMnemonic(mnemonic, passphrase string) (*HDWallet, error) {
    // Validate mnemonic
    if !bip39.IsMnemonicValid(mnemonic) {
        return nil, fmt.Errorf("invalid mnemonic phrase")
    }

    // Recover wallet
    return NewHDWallet(mnemonic, passphrase)
}

// RecoverFromPartialMnemonic attempts recovery with missing words
func (mr *MnemonicRecovery) RecoverFromPartialMnemonic(words []string, missing []int) ([]string, error) {
    var possibleMnemonics []string

    // Brute force missing words (careful with performance)
    var tryWords func(index int, current []string)
    tryWords = func(index int, current []string) {
        if index == len(missing) {
            // Try this combination
            mnemonic := strings.Join(current, " ")
            if bip39.IsMnemonicValid(mnemonic) {
                possibleMnemonics = append(possibleMnemonics, mnemonic)
            }
            return
        }

        // Try each word from wordlist
        for _, word := range mr.wordList {
            current[missing[index]] = word
            tryWords(index+1, current)
        }
    }

    currentWords := make([]string, len(words))
    copy(currentWords, words)
    tryWords(0, currentWords)

    return possibleMnemonics, nil
}

// SecureKeyGenerator ensures proper key generation
type SecureKeyGenerator struct {
    entropySource io.Reader
    validator     KeyValidator
}

type KeyValidator interface {
    ValidateKey(key interface{}) error
}

func (skg *SecureKeyGenerator) GenerateKey(keyType KeyType) (*KeyPair, error) {
    // Generate with entropy check
    entropy := make([]byte, 32)
    if _, err := io.ReadFull(skg.entropySource, entropy); err != nil {
        return nil, fmt.Errorf("insufficient entropy: %w", err)
    }

    // Generate key
    keyPair, err := GenerateKeyPair(keyType)
    if err != nil {
        return nil, err
    }

    // Validate key
    if err := skg.validator.ValidateKey(keyPair.PrivateKey); err != nil {
        return nil, fmt.Errorf("key validation failed: %w", err)
    }

    // Clear entropy from memory
    for i := range entropy {
        entropy[i] = 0
    }

    return keyPair, nil
}

// Memory-safe key handling
func SafeKeyOperation(key []byte, operation func([]byte) error) error {
    // Ensure key is cleared after use
    defer func() {
        for i := range key {
            key[i] = 0
        }
    }()

    return operation(key)
}

// KeyAccessControl manages key access permissions
type KeyAccessControl struct {
    permissions map[string]Permission
    audit       AuditLogger
}

type Permission struct {
    CanSign    bool
    CanExport  bool
    CanRotate  bool
    ValidUntil time.Time
}

func (kac *KeyAccessControl) CheckPermission(keyID, operation, userID string) bool {
    perm, exists := kac.permissions[keyID+":"+userID]
    if !exists {
        kac.audit.Log("Permission denied", keyID, operation, userID)
        return false
    }

    // Check expiry
    if time.Now().After(perm.ValidUntil) {
        kac.audit.Log("Permission expired", keyID, operation, userID)
        return false
    }

    // Check operation
    allowed := false
    switch operation {
    case "sign":
        allowed = perm.CanSign
    case "export":
        allowed = perm.CanExport
    case "rotate":
        allowed = perm.CanRotate
    }

    kac.audit.Log("Permission check", keyID, operation, userID, allowed)
    return allowed
}

func TestKeyManagement(t *testing.T) {
    t.Run("HDWallet", func(t *testing.T) {
        // Generate new wallet
        wallet, err := GenerateNewWallet(24, "test-passphrase")
        require.NoError(t, err)
        require.Len(t, strings.Split(wallet.mnemonic, " "), 24)

        // Derive addresses
        for i := uint32(0); i < 10; i++ {
            key, addr, err := wallet.DeriveAddress(60, 0, 0, i)
            require.NoError(t, err)
            require.NotNil(t, key)
            require.True(t, strings.HasPrefix(addr, "0x"))
        }
    })

    t.Run("KeyEncryption", func(t *testing.T) {
        // Generate key
        key, err := GenerateRandomBytes(32)
        require.NoError(t, err)

        // Encrypt
        encrypted, err := EncryptKey(key, "strong-password")
        require.NoError(t, err)

        // Decrypt
        decrypted, err := DecryptKey(encrypted, "strong-password")
        require.NoError(t, err)
        require.Equal(t, key, decrypted)

        // Wrong password should fail
        _, err = DecryptKey(encrypted, "wrong-password")
        require.Error(t, err)
    })

    t.Run("SecretSharing", func(t *testing.T) {
        ss := &SecretSharing{threshold: 3, shares: 5}

        // Split secret
        secret := []byte("super-secret-key")
        shares, err := ss.SplitSecret(secret)
        require.NoError(t, err)
        require.Len(t, shares, 5)

        // Recover with threshold shares
        recovered, err := ss.RecoverSecret(shares[:3])
        require.NoError(t, err)
        require.Equal(t, secret, recovered)

        // Should fail with insufficient shares
        _, err = ss.RecoverSecret(shares[:2])
        require.Error(t, err)
    })
}