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
	quasar   *quasar.BLS
	finality chan FinalityEvent
	mu       sync.RWMutex

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

	// Cryptography
	certGen *CertificateGenerator
}

// FinalityEvent represents a finalized block with quantum-resistant proofs
type FinalityEvent struct {
	Height    uint64
	BlockID   ids.ID
	Timestamp time.Time
	PQProof   []byte // Post-quantum proof
	BLSProof  []byte // Classical BLS proof
}

// 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
func NewConsensus(params config.Parameters) *ConsensusEngine {
	// Create a simple in-memory store for P-Chain vertices
	// TODO: Use proper persistent storage in production
	store := &memoryStore{}

	return &ConsensusEngine{
		params:    params,
		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
	}

	// For now, use simplified consensus logic
	// In production, this would integrate with quasar's internal methods
	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 aggregate signature
		blsSig, err := e.certGen.GenerateBLSAggregate(blockID, votes)
		if err != nil {
			return fmt.Errorf("failed to generate BLS signature: %w", err)
		}

		// Generate post-quantum certificate
		pqCert, err := e.certGen.GeneratePQCertificate(blockID, votes)
		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.GenerateBLSAggregate(blockID, votes)
		pqCert, _ := testGen.GeneratePQCertificate(blockID, votes)

		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(),
		PQProof:   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 pqProof, blsProof []byte
	if block.Cert != nil {
		pqProof = block.Cert.PQ
		blsProof = block.Cert.BLS
	}
	return FinalityEvent{
		Height:    block.Height,
		BlockID:   blockID,
		Timestamp: block.Timestamp,
		PQProof:   pqProof,
		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(),
	}
}