blueprint.go.tmpl

import (
	"fmt"
	"math/bits"
	"reflect"
	"slices"
	"sync"

	"github.com/consensys/gnark-crypto/ecc"
	"{{ .FieldPackagePath }}"
	"{{ .FieldPackagePath }}/polynomial"
	"github.com/consensys/gnark-crypto/hash"
	"github.com/consensys/gnark/constraint"
	"github.com/consensys/gnark/internal/gkr/gkrcore"
)

func init() {
	// Register GKR blueprint types for CBOR serialization with explicit tag numbers
	constraint.RegisterGkrBlueprintTypes(ecc.{{.FieldID}}, BlueprintSolve{}, BlueprintProve{}, BlueprintGetAssignment{})
}

// circuitEvaluator evaluates all gates in a circuit for one instance
type circuitEvaluator struct {
	evaluators []gateEvaluator // one evaluator per wire
}

// BlueprintSolve is a {{.FieldID}}-specific blueprint for solving GKR circuit instances.
type BlueprintSolve struct {
	// Circuit structure (serialized)
	Circuit     gkrcore.SerializableCircuit
	NbInstances uint32

	// Not serialized - recreated lazily at solve time
	assignments    WireAssignment   `cbor:"-"`
	evaluatorPool  sync.Pool        `cbor:"-"` // pool of circuitEvaluator, lazy-initialized
	outputWires    []int            `cbor:"-"`
	maxOutputLevel constraint.Level `cbor:"-"` // highest output level across all solve instances

	lock sync.Mutex `cbor:"-"`
}

// Ensures BlueprintSolve implements BlueprintStateful
var _ constraint.BlueprintStateful[constraint.U64] = (*BlueprintSolve)(nil)

// Equal returns true if the serialized fields of two BlueprintSolve are equal.
// Used for testing serialization round-trips.
func (b *BlueprintSolve) Equal(other constraint.BlueprintComparable) bool {
	if other == nil {
		return false
	}
	o, ok := other.(*BlueprintSolve)
	if !ok {
		return false
	}
	return b.NbInstances == o.NbInstances && reflect.DeepEqual(b.Circuit, o.Circuit) && reflect.DeepEqual(b.assignments, o.assignments)
}

// Reset implements BlueprintStateful.
// It is used to initialize the blueprint for the current circuit.
func (b *BlueprintSolve) Reset() {
	b.evaluatorPool.New = func() interface{} {
		ce := &circuitEvaluator{
			evaluators: make([]gateEvaluator, len(b.Circuit)),
		}
		for wI := range b.Circuit {
			w := &b.Circuit[wI]
			if !w.IsInput() {
				ce.evaluators[wI] = newGateEvaluator(w.Gate.Evaluate, len(w.Inputs))
			}
		}
		return ce
	}

	b.outputWires = b.Circuit.Outputs()

	assignments := make(WireAssignment, len(b.Circuit))
	nbPaddedInstances := ecc.NextPowerOfTwo(uint64(b.NbInstances))
	for i := range assignments {
		assignments[i] = make(polynomial.MultiLin, nbPaddedInstances)
	}
	b.assignments = assignments
}

// Solve implements the BlueprintStateful interface.
func (b *BlueprintSolve) Solve(s constraint.Solver[constraint.U64], inst constraint.Instruction) error {

	// Get a circuit evaluator from the pool
	ce := b.evaluatorPool.Get().(*circuitEvaluator)
	defer b.evaluatorPool.Put(ce)

	// Format: [0]=totalSize, [1]=instanceIndex, [2...]=input linear expressions
	instanceI := int(inst.Calldata[1])
	calldata := inst.Calldata[2:]

	// The test engine runs the instruction before "finalize" is called. We need to avoid attempting to access an uninitialized slice
	if len(b.assignments[0]) <= instanceI {
		b.lock.Lock()
		defer b.lock.Unlock()
		for wI := range b.assignments {
			for len(b.assignments[wI]) <= instanceI {
				b.assignments[wI] = append(b.assignments[wI], {{ .ElementType }}{})
			}
		}
	}

	// Process all wires in topological order (circuit is already sorted)
	for wI := range b.Circuit {
		w := &b.Circuit[wI]

		if w.IsInput() {
			val, delta := s.Read(calldata)
			calldata = calldata[delta:]
			// Copy directly from constraint.U64 to {{ .ElementType }} (both in Montgomery form)
			copy(b.assignments[wI][instanceI][:], val[:])
		} else {
			// Get evaluator for this wire from the circuit evaluator
			evaluator := &ce.evaluators[wI]

			// Push gate inputs
			for _, inWI := range w.Inputs {
				evaluator.pushInput(b.assignments[inWI][instanceI])
			}

			// Evaluate the gate
			b.assignments[wI][instanceI].Set(evaluator.evaluate())
		}
	}

	// Set output wires (copy {{ .ElementType }} to U64 in Montgomery form)
	for outI, outWI := range b.outputWires {
		var val constraint.U64
		copy(val[:], b.assignments[outWI][instanceI][:])
		s.SetValue(uint32(outI+int(inst.WireOffset)), val)
	}

	return nil
}

// SetNbInstances sets the number of instances for the blueprint
func (b *BlueprintSolve) SetNbInstances(nbInstances uint32) {
	b.NbInstances = nbInstances
}

// CalldataSize implements Blueprint
func (b *BlueprintSolve) CalldataSize() int {
	return -1 // variable size
}

// NbConstraints implements Blueprint
func (b *BlueprintSolve) NbConstraints() int {
	return 0
}

// NbOutputs implements Blueprint
func (b *BlueprintSolve) NbOutputs(inst constraint.Instruction) int {
	if b.outputWires == nil {
		b.outputWires = b.Circuit.Outputs()
	}
	return len(b.outputWires)
}

// UpdateInstructionTree implements Blueprint
func (b *BlueprintSolve) UpdateInstructionTree(inst constraint.Instruction, tree constraint.InstructionTree) constraint.Level {
	maxLevel := constraint.LevelUnset

	// Format: [0]=totalSize, [1]=instanceIndex, [2...]=input linear expressions
	offset := 2 // skip size and instance index

	// Parse input linear expressions
	for offset < len(inst.Calldata) {
		n := int(inst.Calldata[offset]) // number of terms in this linear expression
		offset++

		for range n {
			wireID := inst.Calldata[offset+1]
			offset += 2
			if !tree.HasWire(wireID) {
				continue
			}
			if level := tree.GetWireLevel(wireID); level > maxLevel {
				maxLevel = level
			}
		}
	}

	outputLevel := maxLevel + 1
	for i := range b.outputWires {
		tree.InsertWire(uint32(i+int(inst.WireOffset)), outputLevel)
	}

	b.maxOutputLevel = max(b.maxOutputLevel, outputLevel)

	return outputLevel
}

// BlueprintProve is a {{.FieldID}}-specific blueprint for generating GKR proofs.
type BlueprintProve struct {
	SolveBlueprintID constraint.BlueprintID
	SolveBlueprint   *BlueprintSolve `cbor:"-"` // not serialized, set at compile time
	Schedule         constraint.GkrProvingSchedule
	HashName         string

	lock sync.Mutex
}

// Ensures BlueprintProve implements BlueprintSolvable
var _ constraint.BlueprintSolvable[constraint.U64] = (*BlueprintProve)(nil)

// Equal returns true if the serialized fields of two BlueprintProve are equal.
func (b *BlueprintProve) Equal(other constraint.BlueprintComparable) bool {
	if other == nil {
		return false
	}
	o, ok := other.(*BlueprintProve)
	if !ok {
		return false
	}
	return b.SolveBlueprintID == o.SolveBlueprintID && b.HashName == o.HashName
}

// Solve implements the BlueprintSolvable interface for proving.
func (b *BlueprintProve) Solve(s constraint.Solver[constraint.U64], inst constraint.Instruction) error {
	b.lock.Lock()
	defer b.lock.Unlock()

	// Get solve blueprint from solver by ID
	solveBlueprint := s.GetBlueprint(b.SolveBlueprintID).(*BlueprintSolve)

	// Get assignments from solve blueprint (already in {{ .ElementType }} form)
	assignments := solveBlueprint.assignments
	if len(assignments) == 0 {
		return fmt.Errorf("no assignments available for proving")
	}

	nbPaddedInstances := uint32(ecc.NextPowerOfTwo(uint64(solveBlueprint.NbInstances)))
	if nbPaddedInstances != uint32(len(assignments[0])) { // test engine path
		nbPadding := nbPaddedInstances - solveBlueprint.NbInstances
		for wI := range assignments {
			assignments[wI] = slices.Grow(assignments[wI], int(nbPadding))
			toRepeat := assignments[wI][solveBlueprint.NbInstances-1]
			for range nbPadding {
				assignments[wI] = append(assignments[wI], toRepeat)
			}
		}
	} else {
		for wI := range assignments {
			for i := solveBlueprint.NbInstances; i < nbPaddedInstances; i++ {
				assignments[wI][i] = assignments[wI][i-1]
			}
		}
	}

	// Create hasher and write base challenges
	hsh := hash.NewHash(b.HashName + "_{{.FieldID}}")

	// Read initial challenges from instruction calldata (parse dynamically, no metadata)
	// Format: [0]=totalSize, [1...]=challenge linear expressions
	calldata := inst.Calldata[1:] // skip size prefix
	for len(calldata) != 0 {
		val, delta := s.Read(calldata)
		calldata = calldata[delta:]

		// Copy directly from constraint.U64 to {{ .ElementType }} (both in Montgomery form)
		var challenge {{ .ElementType }}
		copy(challenge[:], val[:])
		challengeBytes := challenge.Bytes()
		hsh.Write(challengeBytes[:])
	}

	// Call the {{.FieldID}}-specific Prove function (assignments already WireAssignment type)
	proof, err := Prove(solveBlueprint.Circuit, b.Schedule, assignments, hsh)
	if err != nil {
		return fmt.Errorf("{{.FieldID}} prove failed: %w", err)
	}

	for i, elem := range proof.flatten() {
		var val constraint.U64
		copy(val[:], (*elem)[:])
		s.SetValue(uint32(i+int(inst.WireOffset)), val)
	}

	return nil
}

// CalldataSize implements Blueprint
func (b *BlueprintProve) CalldataSize() int {
	return -1 // variable size
}

// NbConstraints implements Blueprint
func (b *BlueprintProve) NbConstraints() int {
	return 0
}

func (b *BlueprintProve) proofSize() int {
	if b.SolveBlueprint.NbInstances < 2 {
		return 0
	}
	nbPaddedInstances := ecc.NextPowerOfTwo(uint64(b.SolveBlueprint.NbInstances))
	logNbInstances := bits.TrailingZeros64(nbPaddedInstances)
	return b.SolveBlueprint.Circuit.ProofSize(b.Schedule, logNbInstances)
}

// NbOutputs implements Blueprint
func (b *BlueprintProve) NbOutputs(inst constraint.Instruction) int {
	return b.proofSize()
}

// UpdateInstructionTree implements Blueprint
func (b *BlueprintProve) UpdateInstructionTree(inst constraint.Instruction, tree constraint.InstructionTree) constraint.Level {
	maxLevel := b.SolveBlueprint.maxOutputLevel

	// Format: [0]=totalSize, [1...]=challenge linear expressions
	offset := 1 // skip size prefix

	// Parse all challenge linear expressions
	for offset < len(inst.Calldata) {
		n := int(inst.Calldata[offset]) // number of terms in this linear expression
		offset++

		for range n {
			wireID := inst.Calldata[offset+1]
			offset += 2
			if !tree.HasWire(wireID) {
				continue
			}

			maxLevel = max(maxLevel, tree.GetWireLevel(wireID))
		}
	}

	outputLevel := maxLevel + 1
	// Compute proof size from blueprint state
	if b.SolveBlueprint.NbInstances > 0 {
		proofSize := b.proofSize()
		for i := range proofSize {
			tree.InsertWire(uint32(i+int(inst.WireOffset)), outputLevel)
		}
	}

	return outputLevel
}

// BlueprintGetAssignment is a {{.FieldID}}-specific blueprint for retrieving wire assignments.
type BlueprintGetAssignment struct {
	SolveBlueprintID constraint.BlueprintID

	lock sync.Mutex
}

// Ensures BlueprintGetAssignment implements BlueprintSolvable
var _ constraint.BlueprintSolvable[constraint.U64] = (*BlueprintGetAssignment)(nil)

// Equal returns true if the serialized fields of two BlueprintGetAssignment are equal.
func (b *BlueprintGetAssignment) Equal(other constraint.BlueprintComparable) bool {
	if other == nil {
		return false
	}
	o, ok := other.(*BlueprintGetAssignment)
	if !ok {
		return false
	}
	return b.SolveBlueprintID == o.SolveBlueprintID
}

// Solve implements the BlueprintSolvable interface for getting assignments.
func (b *BlueprintGetAssignment) Solve(s constraint.Solver[constraint.U64], inst constraint.Instruction) error {
	b.lock.Lock()
	defer b.lock.Unlock()

	// Read wireI and instanceI from calldata
	// Format: [0]=totalSize, [1]=wireI, [2]=instanceI, [3...]=dependency linear expression
	wireI := int(inst.Calldata[1])
	instanceI := int(inst.Calldata[2])

	// Get solve blueprint from solver by ID
	solveBlueprint := s.GetBlueprint(b.SolveBlueprintID).(*BlueprintSolve)

	var v constraint.U64
	copy(v[:], solveBlueprint.assignments[wireI][instanceI][:])
	// Set output wire
	s.SetValue(inst.WireOffset, v)
	return nil
}

// CalldataSize implements Blueprint
func (b *BlueprintGetAssignment) CalldataSize() int {
	return -1 // variable size: [wireI, instanceI, dependency_linear_expression]
}

// NbConstraints implements Blueprint
func (b *BlueprintGetAssignment) NbConstraints() int {
	return 0
}

// NbOutputs implements Blueprint
func (b *BlueprintGetAssignment) NbOutputs(inst constraint.Instruction) int {
	return 1 // returns one assignment value
}

// UpdateInstructionTree implements Blueprint
func (b *BlueprintGetAssignment) UpdateInstructionTree(inst constraint.Instruction, tree constraint.InstructionTree) constraint.Level {
	maxLevel := constraint.LevelUnset

	// Format: [0]=totalSize, [1]=wireI, [2]=instanceI, [3...]=dependency linear expression
	offset := 3 // skip size, wireI, and instanceI

	// Parse dependency linear expression
	// This ensures we run after the solve instruction for this instance
	n := int(inst.Calldata[offset])
	offset++

	for range n {
		wireID := inst.Calldata[offset+1]
		offset += 2
		if !tree.HasWire(wireID) {
			continue
		}
		if level := tree.GetWireLevel(wireID); level > maxLevel {
			maxLevel = level
		}
	}

	outputLevel := maxLevel + 1
	tree.InsertWire(inst.WireOffset, outputLevel)
	return outputLevel
}

// NewBlueprints creates and registers all GKR blueprints for {{.FieldID}}
func NewBlueprints(circuit gkrcore.SerializableCircuit, schedule constraint.GkrProvingSchedule, hashName string, compiler constraint.CustomizableSystem) gkrcore.Blueprints {
	// Create and register solve blueprint
	solve := &BlueprintSolve{Circuit: circuit}
	solveID := compiler.AddBlueprint(solve)

	// Create and register prove blueprint
	prove := &BlueprintProve{
		SolveBlueprintID: solveID,
		SolveBlueprint:   solve,
		Schedule:         schedule,
		HashName:         hashName,
	}
	proveID := compiler.AddBlueprint(prove)

	// Create and register GetAssignment blueprint
	getAssignment := &BlueprintGetAssignment{
		SolveBlueprintID: solveID,
	}
	getAssignmentID := compiler.AddBlueprint(getAssignment)

	return gkrcore.Blueprints{
		SolveID:         solveID,
		Solve:           solve,
		ProveID:         proveID,
		GetAssignmentID: getAssignmentID,
	}
}