// Package pqid implements hybrid post-quantum player identity.
//
// A player's key is two keys: classical Ed25519 and post-quantum ML-DSA-65
// (NIST FIPS 204). Both signatures are required for every authentication, so
// the identity holds if *either* algorithm survives:
//
//   - Ed25519 alone falls to Shor's algorithm on a sufficiently large quantum
//     computer.
//   - ML-DSA is young. Lattice cryptanalysis is an active field, and betting
//     everything on a 2024 standard would be its own kind of naive.
//
// Requiring both is the composition NIST and the IETF recommend: an attacker
// must break lattice assumptions *and* elliptic curves, not either one.
//
// What this does and does not protect:
//
//	protected     player identity, session authentication, bet authorisation
//	protected     round fairness — commit-reveal is SHA-256/HMAC, and Grover
//	              only halves the security margin, leaving ~128 bits
//	NOT protected Bitcoin and Lightning settlement, which sign with secp256k1.
//	              No application-layer choice can change that.
package pqid

import (
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"

	"github.com/cloudflare/circl/sign/mldsa/mldsa65"
)

// Sizes of the wire encodings.
const (
	EdPublicKeySize = ed25519.PublicKeySize // 32
	EdSignatureSize = ed25519.SignatureSize // 64

	PQPublicKeySize = mldsa65.PublicKeySize // 1952
	PQSignatureSize = mldsa65.SignatureSize // 3309

	PublicKeySize = EdPublicKeySize + PQPublicKeySize
	SignatureSize = EdSignatureSize + PQSignatureSize

	// IDSize is the length of the account identifier derived from a key.
	IDSize = sha256.Size
)

// context binds signatures to this application, so a signature produced for
// some other ML-DSA protocol cannot be replayed here.
var context = []byte("quantum-arcade/v1")

var (
	ErrMalformedKey       = errors.New("pqid: malformed public key")
	ErrMalformedSignature = errors.New("pqid: malformed signature")
	ErrClassicalFailed    = errors.New("pqid: ed25519 signature did not verify")
	ErrPostQuantumFailed  = errors.New("pqid: ML-DSA signature did not verify")
)

// PublicKey is a player's hybrid identity.
type PublicKey struct {
	Ed ed25519.PublicKey
	PQ *mldsa65.PublicKey
}

// PrivateKey is the matching secret half. Clients hold this; the server never
// sees it.
type PrivateKey struct {
	Ed ed25519.PrivateKey
	PQ *mldsa65.PrivateKey
}

// GenerateKey creates a hybrid keypair.
func GenerateKey(rand io.Reader) (*PublicKey, *PrivateKey, error) {
	edPub, edPriv, err := ed25519.GenerateKey(rand)
	if err != nil {
		return nil, nil, fmt.Errorf("pqid: generating ed25519 key: %w", err)
	}
	pqPub, pqPriv, err := mldsa65.GenerateKey(rand)
	if err != nil {
		return nil, nil, fmt.Errorf("pqid: generating ML-DSA key: %w", err)
	}
	return &PublicKey{Ed: edPub, PQ: pqPub}, &PrivateKey{Ed: edPriv, PQ: pqPriv}, nil
}

// Bytes encodes the public key as ed25519 || ML-DSA.
func (p *PublicKey) Bytes() []byte {
	out := make([]byte, 0, PublicKeySize)
	out = append(out, p.Ed...)
	pq, err := p.PQ.MarshalBinary()
	if err != nil {
		// MarshalBinary on a valid key cannot fail; a failure here means the
		// key is corrupt, and silently returning a short key would be worse.
		panic("pqid: marshalling ML-DSA public key: " + err.Error())
	}
	return append(out, pq...)
}

// Hex renders the public key for transport.
func (p *PublicKey) Hex() string { return hex.EncodeToString(p.Bytes()) }

// ID is the account identifier: SHA-256 over the whole hybrid key.
//
// The ledger keys on this rather than the raw key because the hybrid key is
// nearly 2KB, and a fixed 32-byte identifier keeps indexes small. Hashing also
// means the identifier is stable in length no matter how the key evolves.
func (p *PublicKey) ID() []byte {
	sum := sha256.Sum256(p.Bytes())
	return sum[:]
}

// ParsePublicKey decodes a hybrid public key from its wire encoding.
func ParsePublicKey(b []byte) (*PublicKey, error) {
	if len(b) != PublicKeySize {
		return nil, fmt.Errorf("%w: got %d bytes, want %d",
			ErrMalformedKey, len(b), PublicKeySize)
	}
	ed := ed25519.PublicKey(append([]byte(nil), b[:EdPublicKeySize]...))

	var pq mldsa65.PublicKey
	if err := pq.UnmarshalBinary(b[EdPublicKeySize:]); err != nil {
		return nil, fmt.Errorf("%w: %v", ErrMalformedKey, err)
	}
	return &PublicKey{Ed: ed, PQ: &pq}, nil
}

// ParsePublicKeyHex decodes a hex-encoded hybrid public key.
func ParsePublicKeyHex(s string) (*PublicKey, error) {
	b, err := hex.DecodeString(s)
	if err != nil {
		return nil, fmt.Errorf("%w: not hex", ErrMalformedKey)
	}
	return ParsePublicKey(b)
}

// Sign produces both signatures over the message, concatenated.
func Sign(priv *PrivateKey, msg []byte) ([]byte, error) {
	edSig := ed25519.Sign(priv.Ed, msg)

	pqSig := make([]byte, PQSignatureSize)
	// randomized=false gives deterministic (hedged) signatures, so a bad RNG
	// on a phone cannot leak the key through signature randomness.
	if err := mldsa65.SignTo(priv.PQ, msg, context, false, pqSig); err != nil {
		return nil, fmt.Errorf("pqid: ML-DSA signing: %w", err)
	}

	out := make([]byte, 0, SignatureSize)
	out = append(out, edSig...)
	return append(out, pqSig...), nil
}

// Verify checks both signatures. Both must pass.
//
// It deliberately reports which half failed: a mismatch between the two is
// diagnostic — it means a client is half-upgraded or something is tampering
// with one algorithm — and that is worth surfacing rather than flattening into
// a generic failure. Nothing secret is revealed by saying which one broke.
func Verify(pub *PublicKey, msg, sig []byte) error {
	if len(sig) != SignatureSize {
		return fmt.Errorf("%w: got %d bytes, want %d",
			ErrMalformedSignature, len(sig), SignatureSize)
	}
	if !ed25519.Verify(pub.Ed, msg, sig[:EdSignatureSize]) {
		return ErrClassicalFailed
	}
	if !mldsa65.Verify(pub.PQ, msg, context, sig[EdSignatureSize:]) {
		return ErrPostQuantumFailed
	}
	return nil
}

// PrivateFromBytes reconstructs a private key from stored material.
//
// The Ed25519 half is stored as its 32-byte seed rather than the expanded
// key, because the seed is the canonical form and cannot encode an
// inconsistent pair. The ML-DSA half is stored in its own binary encoding.
func PrivateFromBytes(edSeed, pqKey []byte) (*PrivateKey, error) {
	if len(edSeed) != ed25519.SeedSize {
		return nil, fmt.Errorf("%w: ed25519 seed is %d bytes, want %d",
			ErrMalformedKey, len(edSeed), ed25519.SeedSize)
	}
	var pq mldsa65.PrivateKey
	if err := pq.UnmarshalBinary(pqKey); err != nil {
		return nil, fmt.Errorf("%w: ML-DSA private key: %v", ErrMalformedKey, err)
	}
	return &PrivateKey{
		Ed: ed25519.NewKeyFromSeed(edSeed),
		PQ: &pq,
	}, nil
}

// PublicFromPrivate derives the public half.
//
// Deriving rather than storing means a client cannot present a public key that
// does not match the key it signs with — a mismatch that would otherwise only
// surface as a confusing authentication failure.
func PublicFromPrivate(priv *PrivateKey) (*PublicKey, error) {
	edPub, ok := priv.Ed.Public().(ed25519.PublicKey)
	if !ok {
		return nil, fmt.Errorf("%w: ed25519 private key has no public half", ErrMalformedKey)
	}
	pqPub, ok := priv.PQ.Public().(*mldsa65.PublicKey)
	if !ok {
		return nil, fmt.Errorf("%w: ML-DSA private key has no public half", ErrMalformedKey)
	}
	return &PublicKey{Ed: edPub, PQ: pqPub}, nil
}