consensus/engine/pq/consensus.go

// Package pq implements the post-quantum consensus engine that combines
// classical and quantum-resistant consensus mechanisms.
package pq

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/luxfi/consensus/config"
	"github.com/luxfi/consensus/core/dag"
	"github.com/luxfi/consensus/protocol/quasar"
	"github.com/luxfi/ids"
)

// ConsensusEngine implements post-quantum consensus combining classical and quantum-resistant security
type ConsensusEngine struct {
	params   config.Parameters
	pqMode   config.PQMode
	quasar   *quasar.BLS
	finality chan FinalityEvent
	mu       sync.RWMutex

	// State tracking
	height    uint64
	round     uint64
	finalized map[ids.ID]bool

	// Cryptography
	certGen *CertificateGenerator
}

// PQMode returns the post-quantum signature mode this engine was built with.
func (e *ConsensusEngine) PQMode() config.PQMode { return e.pqMode }

// FinalityEvent represents a finalized block with quantum-resistant proofs.
// Carries the configured layers for parallel verification: classical
// BLS fast-path, Corona (Ring-LWE) threshold, optional Pulsar (Module-
// LWE) threshold, and the per-validator ML-DSA-65 identity rollup.
// Defence in depth across one classical curve and (one or two) PQ
// lattice families.
type FinalityEvent struct {
	Height      uint64
	BlockID     ids.ID
	Timestamp   time.Time
	BLSProof    []byte // Classical BLS-12-381 aggregate (optional)
	Corona      []byte // Corona (Ring-LWE) threshold sig
	Pulsar      []byte // Pulsar-M (Module-LWE) threshold sig (optional)
	MLDSARollup []byte // Per-validator ML-DSA-65 rolled up via STARK/Groth16
}

// memoryStore is a simple in-memory implementation of dag.Store for P-Chain vertices
type memoryStore struct {
	vertices map[quasar.VertexID]*memoryVertex
	heads    []quasar.VertexID
	mu       sync.RWMutex
}

type memoryVertex struct {
	id       quasar.VertexID
	parents  []quasar.VertexID
	author   string
	round    uint64
	children []quasar.VertexID
}

func (v *memoryVertex) ID() quasar.VertexID        { return v.id }
func (v *memoryVertex) Parents() []quasar.VertexID { return v.parents }
func (v *memoryVertex) Author() string             { return v.author }
func (v *memoryVertex) Round() uint64              { return v.round }

func (s *memoryStore) Head() []quasar.VertexID {
	if s.vertices == nil {
		s.vertices = make(map[quasar.VertexID]*memoryVertex)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]quasar.VertexID, len(s.heads))
	copy(result, s.heads)
	return result
}

func (s *memoryStore) Get(id quasar.VertexID) (dag.BlockView[quasar.VertexID], bool) {
	if s.vertices == nil {
		s.vertices = make(map[quasar.VertexID]*memoryVertex)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	vertex, exists := s.vertices[id]
	return vertex, exists
}

func (s *memoryStore) Children(id quasar.VertexID) []quasar.VertexID {
	if s.vertices == nil {
		s.vertices = make(map[quasar.VertexID]*memoryVertex)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	if vertex, exists := s.vertices[id]; exists {
		result := make([]quasar.VertexID, len(vertex.children))
		copy(result, vertex.children)
		return result
	}
	return []quasar.VertexID{}
}

// NewConsensus creates a new post-quantum consensus engine.
// Uses in-memory storage suitable for single-node operation.
// For distributed deployment, inject a persistent Store implementation.
//
// The engine honours params.PQMode (and the CONSENSUS_PQ_MODE env
// override) to select which signature layers are active. Only the BLS
// fast path is wired today; PQ paths are ready for the signer hookup.
func NewConsensus(params config.Parameters) *ConsensusEngine {
	store := &memoryStore{}

	mode, _ := config.PQModeFromEnv(params.PQMode)

	return &ConsensusEngine{
		params:    params,
		pqMode:    mode,
		quasar:    quasar.NewBLS(params, store),
		finality:  make(chan FinalityEvent, 100),
		finalized: make(map[ids.ID]bool),
	}
}

// Initialize sets up the PQ engine with keys
func (e *ConsensusEngine) Initialize(ctx context.Context, blsKey, pqKey []byte) error {
	e.certGen = NewCertificateGenerator(blsKey, pqKey)
	return e.quasar.Initialize(ctx, blsKey, pqKey)
}

// ProcessBlock processes a block through PQ consensus
func (e *ConsensusEngine) ProcessBlock(ctx context.Context, blockID ids.ID, votes map[string]int) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	// Check if already finalized
	if e.finalized[blockID] {
		return nil
	}

	// Simplified majority voting - the quasar.BLS handles DAG finality separately
	totalVotes := 0
	maxVotes := 0
	bestBlock := ""

	for block, count := range votes {
		totalVotes += count
		if count > maxVotes {
			maxVotes = count
			bestBlock = block
		}
	}

	// Check if we have enough votes for the best block
	if totalVotes == 0 || float64(maxVotes)/float64(totalVotes) < e.params.Alpha {
		return fmt.Errorf("insufficient votes for finality: %d/%d for block %s", maxVotes, totalVotes, bestBlock)
	}

	// Generate real certificates using cryptography
	var cert *quasar.CertBundle
	if e.certGen != nil {
		// Generate BLS signature for this block
		blsSig, err := e.certGen.SignBlock(blockID)
		if err != nil {
			return fmt.Errorf("failed to generate BLS signature: %w", err)
		}

		// Generate post-quantum certificate
		pqCert, err := e.certGen.GeneratePQSignature(blockID)
		if err != nil {
			return fmt.Errorf("failed to generate PQ certificate: %w", err)
		}

		cert = &quasar.CertBundle{
			BLSAgg: blsSig,
			PQCert: pqCert,
		}
	} else {
		// Fallback to test keys if not initialized
		blsKey, pqKey := GenerateTestKeys()
		testGen := NewCertificateGenerator(blsKey, pqKey)

		blsSig, _ := testGen.SignBlock(blockID)
		pqCert, _ := testGen.GeneratePQSignature(blockID)

		cert = &quasar.CertBundle{
			BLSAgg: blsSig,
			PQCert: pqCert,
		}
	}

	// Validate certificate
	if len(cert.BLSAgg) == 0 || len(cert.PQCert) == 0 {
		return fmt.Errorf("failed to generate valid certificates")
	}

	// Mark as finalized
	e.finalized[blockID] = true
	e.height++

	// Emit finality event
	select {
	case e.finality <- FinalityEvent{
		Height:      e.height,
		BlockID:     blockID,
		Timestamp:   time.Now(),
		MLDSARollup: cert.PQCert,
		BLSProof:    cert.BLSAgg,
	}:
	case <-ctx.Done():
		return ctx.Err()
	default:
		// Channel full, drop oldest
	}

	return nil
}

// IsFinalized checks if a block has achieved PQ finality
func (e *ConsensusEngine) IsFinalized(blockID ids.ID) bool {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.finalized[blockID]
}

// FinalityChannel returns the channel for finality events
func (e *ConsensusEngine) FinalityChannel() <-chan FinalityEvent {
	return e.finality
}

// Height returns the current finalized height
func (e *ConsensusEngine) Height() uint64 {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.height
}

// blockToFinalityEvent converts a quasar.Block to a FinalityEvent
func blockToFinalityEvent(block *quasar.Block) FinalityEvent {
	blockID, _ := ids.FromString(block.Hash)
	var mldsaRollup, corona, pulsar, blsProof []byte
	if block.Cert != nil {
		mldsaRollup = block.Cert.MLDSARollup
		corona = block.Cert.Corona
		pulsar = block.Cert.Pulsar
		blsProof = block.Cert.BLS
	}
	return FinalityEvent{
		Height:      block.Height,
		BlockID:     blockID,
		Timestamp:   block.Timestamp,
		MLDSARollup: mldsaRollup,
		Corona:      corona,
		Pulsar:      pulsar,
		BLSProof:    blsProof,
	}
}

// SetFinalizedCallback sets a callback for finalized blocks
func (e *ConsensusEngine) SetFinalizedCallback(cb func(FinalityEvent)) {
	e.quasar.SetFinalizedCallback(func(block *quasar.Block) {
		cb(blockToFinalityEvent(block))
	})
}

// Metrics returns engine metrics
func (e *ConsensusEngine) Metrics() map[string]interface{} {
	e.mu.RLock()
	defer e.mu.RUnlock()

	return map[string]interface{}{
		"height":     e.height,
		"round":      e.round,
		"finalized":  len(e.finalized),
		"k":          e.params.K,
		"alpha":      e.params.Alpha,
		"beta":       e.params.Beta,
		"block_time": e.params.BlockTime.String(),
	}
}