package main

import (
    "crypto/rand"
    "fmt"
    "log"
    "github.com/luxfi/crypto/mldsa"
)

func generateMLDSAKeys() {
    // Generate ML-DSA-65 key pair (recommended for most applications)
    publicKey, privateKey, err := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
    if err != nil {
        log.Fatal("Key generation failed:", err)
    }

    // Serialize keys for storage
    pubBytes, _ := publicKey.MarshalBinary()
    privBytes, _ := privateKey.MarshalBinary()

    fmt.Printf("ML-DSA-65 Public Key Size: %d bytes\n", len(pubBytes))
    fmt.Printf("ML-DSA-65 Private Key Size: %d bytes\n", len(privBytes))

    // Generate different security levels
    modes := []mldsa.Mode{mldsa.MLDSA44, mldsa.MLDSA65, mldsa.MLDSA87}
    for _, mode := range modes {
        pk, sk, _ := mldsa.GenerateKey(rand.Reader, mode)
        fmt.Printf("Mode %v: PK=%d bytes, SK=%d bytes\n",
            mode, pk.Size(), sk.Size())
    }
}

func signWithMLDSA(privateKey *mldsa.PrivateKey, message []byte) ([]byte, error) {
    // Sign the message
    signature, err := privateKey.Sign(rand.Reader, message, nil)
    if err != nil {
        return nil, fmt.Errorf("signing failed: %w", err)
    }

    fmt.Printf("Signature size: %d bytes\n", len(signature))
    return signature, nil
}

func verifyMLDSA(publicKey *mldsa.PublicKey, message, signature []byte) bool {
    // Verify the signature
    err := publicKey.Verify(message, signature)
    return err == nil
}

// Example usage
func exampleMLDSA() {
    message := []byte("Quantum-resistant signature test")

    // Generate keys
    publicKey, privateKey, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)

    // Sign
    signature, err := signWithMLDSA(privateKey, message)
    if err != nil {
        log.Fatal(err)
    }

    // Verify
    valid := verifyMLDSA(publicKey, message, signature)
    fmt.Printf("Signature valid: %v\n", valid)
}

package main

import (
    "crypto/rand"
    "fmt"
    "github.com/luxfi/crypto/mlkem"
)

func quantumSafeKeyExchange() {
    // Alice generates key pair
    alicePublic, alicePrivate, err := mlkem.GenerateKey(rand.Reader, mlkem.MLKEM768)
    if err != nil {
        panic(err)
    }

    // Bob encapsulates shared secret using Alice's public key
    ciphertext, bobSharedSecret, err := mlkem.Encapsulate(alicePublic)
    if err != nil {
        panic(err)
    }

    // Alice decapsulates to get the same shared secret
    aliceSharedSecret, err := mlkem.Decapsulate(alicePrivate, ciphertext)
    if err != nil {
        panic(err)
    }

    // Verify both parties have the same secret
    if bytes.Equal(aliceSharedSecret, bobSharedSecret) {
        fmt.Println("Key exchange successful!")
        fmt.Printf("Shared secret: %x\n", aliceSharedSecret)
    }
}

type HybridKeyExchange struct {
    mlkemPublic  *mlkem.PublicKey
    mlkemPrivate *mlkem.PrivateKey
    ecdhPublic   *ecdh.PublicKey
    ecdhPrivate  *ecdh.PrivateKey
}

func (hke *HybridKeyExchange) GenerateKeys() error {
    // Generate ML-KEM keys
    mlkemPub, mlkemPriv, err := mlkem.GenerateKey(rand.Reader, mlkem.MLKEM768)
    if err != nil {
        return err
    }
    hke.mlkemPublic = mlkemPub
    hke.mlkemPrivate = mlkemPriv

    // Generate ECDH keys (X25519)
    ecdhPriv, err := ecdh.X25519().GenerateKey(rand.Reader)
    if err != nil {
        return err
    }
    hke.ecdhPrivate = ecdhPriv
    hke.ecdhPublic = ecdhPriv.PublicKey()

    return nil
}

func (hke *HybridKeyExchange) Encapsulate(peerMLKEM *mlkem.PublicKey, peerECDH *ecdh.PublicKey) ([]byte, []byte, error) {
    // ML-KEM encapsulation
    mlkemCiphertext, mlkemSecret, err := mlkem.Encapsulate(peerMLKEM)
    if err != nil {
        return nil, nil, err
    }

    // ECDH shared secret
    ecdhSecret, err := hke.ecdhPrivate.ECDH(peerECDH)
    if err != nil {
        return nil, nil, err
    }

    // Combine secrets using KDF
    combinedSecret := kdf.DeriveKey(
        append(mlkemSecret, ecdhSecret...),
        []byte("hybrid-key-exchange"),
        32,
    )

    // Return ciphertext and combined secret
    hybridCiphertext := append(mlkemCiphertext, hke.ecdhPublic.Bytes()...)
    return hybridCiphertext, combinedSecret, nil
}

package main

import (
    "crypto/rand"
    "fmt"
    "github.com/luxfi/crypto/slhdsa"
)

func hashBasedSignatures() {
    // Generate SLH-DSA key pair
    publicKey, privateKey, err := slhdsa.GenerateKey(rand.Reader, slhdsa.SLHDSA_SHA2_128f)
    if err != nil {
        panic(err)
    }

    message := []byte("Hash-based signature test")

    // Sign message
    signature, err := privateKey.Sign(rand.Reader, message, nil)
    if err != nil {
        panic(err)
    }

    fmt.Printf("SLH-DSA signature size: %d bytes\n", len(signature))

    // Verify signature
    err = publicKey.Verify(message, signature)
    if err == nil {
        fmt.Println("Signature verified successfully")
    }
}

// Demonstrate small key, large signature trade-off
func compareSignatureSizes() {
    modes := []struct {
        name string
        mode slhdsa.Mode
    }{
        {"SLH-DSA-128f", slhdsa.SLHDSA_SHA2_128f},
        {"SLH-DSA-192f", slhdsa.SLHDSA_SHA2_192f},
        {"SLH-DSA-256f", slhdsa.SLHDSA_SHA2_256f},
    }

    for _, m := range modes {
        pk, sk, _ := slhdsa.GenerateKey(rand.Reader, m.mode)
        sig, _ := sk.Sign(rand.Reader, []byte("test"), nil)

        pkBytes, _ := pk.MarshalBinary()
        skBytes, _ := sk.MarshalBinary()

        fmt.Printf("%s: PK=%d, SK=%d, Sig=%d bytes\n",
            m.name, len(pkBytes), len(skBytes), len(sig))
    }
}

type HybridSignature struct {
    Classical  []byte // ECDSA or BLS signature
    PostQuantum []byte // ML-DSA signature
}

type HybridSigner struct {
    classicalKey  crypto.Signer       // ECDSA or BLS
    quantumKey    *mldsa.PrivateKey   // ML-DSA
}

func (hs *HybridSigner) Sign(message []byte) (*HybridSignature, error) {
    // Sign with classical algorithm
    classicalSig, err := hs.classicalKey.Sign(rand.Reader, message, crypto.SHA256)
    if err != nil {
        return nil, fmt.Errorf("classical signing failed: %w", err)
    }

    // Sign with post-quantum algorithm
    quantumSig, err := hs.quantumKey.Sign(rand.Reader, message, nil)
    if err != nil {
        return nil, fmt.Errorf("quantum signing failed: %w", err)
    }

    return &HybridSignature{
        Classical:   classicalSig,
        PostQuantum: quantumSig,
    }, nil
}

func (hs *HybridSignature) Verify(message []byte, classicalPub, quantumPub interface{}) bool {
    // Both signatures must be valid
    classicalValid := verifyClassical(classicalPub, message, hs.Classical)
    quantumValid := verifyQuantum(quantumPub, message, hs.PostQuantum)

    return classicalValid && quantumValid
}

type CryptoEra int

const (
    Classical CryptoEra = iota  // Current: ECDSA/BLS only
    Hybrid                      // Transitional: Both algorithms
    PostQuantum                 // Future: PQ algorithms only
)

type CryptoConfig struct {
    Era              CryptoEra
    RequireClassical bool
    RequireQuantum   bool
    AcceptLegacy     bool
}

func GetCryptoConfig(timestamp time.Time) CryptoConfig {
    quantumTransition := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
    quantumOnly := time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC)

    switch {
    case timestamp.Before(quantumTransition):
        return CryptoConfig{
            Era:              Classical,
            RequireClassical: true,
            RequireQuantum:   false,
            AcceptLegacy:     true,
        }
    case timestamp.Before(quantumOnly):
        return CryptoConfig{
            Era:              Hybrid,
            RequireClassical: true,
            RequireQuantum:   true,
            AcceptLegacy:     true,
        }
    default:
        return CryptoConfig{
            Era:              PostQuantum,
            RequireClassical: false,
            RequireQuantum:   true,
            AcceptLegacy:     false,
        }
    }
}

func batchVerifyMLDSA(publicKeys []*mldsa.PublicKey, messages [][]byte, signatures [][]byte) bool {
    if len(publicKeys) != len(messages) || len(messages) != len(signatures) {
        return false
    }

    // Parallel verification using goroutines
    results := make(chan bool, len(publicKeys))

    for i := range publicKeys {
        go func(idx int) {
            err := publicKeys[idx].Verify(messages[idx], signatures[idx])
            results <- (err == nil)
        }(i)
    }

    // Collect results
    for i := 0; i < len(publicKeys); i++ {
        if !<-results {
            return false
        }
    }

    return true
}

type PQSignatureCache struct {
    cache *lru.Cache
    mu    sync.RWMutex
}

func (c *PQSignatureCache) VerifyWithCache(pubKey *mldsa.PublicKey, message, signature []byte) bool {
    // Generate cache key
    key := sha256.Sum256(append(append(pubKey.Bytes(), message...), signature...))

    // Check cache
    c.mu.RLock()
    if result, ok := c.cache.Get(key); ok {
        c.mu.RUnlock()
        return result.(bool)
    }
    c.mu.RUnlock()

    // Verify and cache result
    valid := (pubKey.Verify(message, signature) == nil)

    c.mu.Lock()
    c.cache.Add(key, valid)
    c.mu.Unlock()

    return valid
}

func TestPostQuantumInterop(t *testing.T) {
    // Test vector from NIST
    testVectors := []struct {
        algorithm string
        pubKey    string
        message   string
        signature string
        valid     bool
    }{
        // Add NIST test vectors here
    }

    for _, tv := range testVectors {
        t.Run(tv.algorithm, func(t *testing.T) {
            // Decode test vector
            pubKey, _ := hex.DecodeString(tv.pubKey)
            message, _ := hex.DecodeString(tv.message)
            signature, _ := hex.DecodeString(tv.signature)

            // Verify signature
            var valid bool
            switch tv.algorithm {
            case "ML-DSA-65":
                pk, _ := mldsa.ParsePublicKey(pubKey)
                valid = (pk.Verify(message, signature) == nil)
            case "ML-KEM-768":
                // KEM doesn't have signatures, skip
                t.Skip("KEM is for key exchange, not signatures")
            case "SLH-DSA-128f":
                pk, _ := slhdsa.ParsePublicKey(pubKey)
                valid = (pk.Verify(message, signature) == nil)
            }

            require.Equal(t, tv.valid, valid)
        })
    }
}