// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package builder

import (
	"context"
	"sync"
	"time"

	"github.com/luxfi/concurrent/lock"
	"github.com/luxfi/container/linked"
	"github.com/luxfi/database/versiondb"
	"github.com/luxfi/ids"
	"github.com/luxfi/node/vms/example/xsvm/chain"
	"github.com/luxfi/node/vms/example/xsvm/execute"
	"github.com/luxfi/node/vms/example/xsvm/tx"
	"github.com/luxfi/runtime"
	vmcore "github.com/luxfi/vm"

	xsblock "github.com/luxfi/node/vms/example/xsvm/block"
)

const MaxTxsPerBlock = 10

var _ Builder = (*builder)(nil)

type Builder interface {
	SetPreference(preferred ids.ID)
	AddTx(ctx context.Context, tx *tx.Tx) error
	WaitForEvent(ctx context.Context) (vmcore.Message, error)
	BuildBlock(ctx context.Context, blockContext *runtime.Runtime) (chain.Block, error)
}

type builder struct {
	chainRuntime *runtime.Runtime
	chain        chain.Chain

	preference ids.ID
	// pendingTxsCond is awoken once there is at least one pending transaction.
	pendingTxsCond *lock.Cond
	pendingTxs     *linked.Hashmap[ids.ID, *tx.Tx]
}

func New(chainRuntime *runtime.Runtime, chain chain.Chain) Builder {
	return &builder{
		chainRuntime:   chainRuntime,
		chain:          chain,
		preference:     chain.LastAccepted(),
		pendingTxsCond: lock.NewCond(&sync.Mutex{}),
		pendingTxs:     linked.NewHashmap[ids.ID, *tx.Tx](),
	}
}

func (b *builder) SetPreference(preferred ids.ID) {
	b.preference = preferred
}

func (b *builder) AddTx(_ context.Context, newTx *tx.Tx) error {
	txID, err := newTx.ID()
	if err != nil {
		return err
	}

	b.pendingTxsCond.L.Lock()
	defer b.pendingTxsCond.L.Unlock()

	b.pendingTxs.Put(txID, newTx)
	b.pendingTxsCond.Broadcast()
	return nil
}

func (b *builder) WaitForEvent(ctx context.Context) (vmcore.Message, error) {
	b.pendingTxsCond.L.Lock()
	defer b.pendingTxsCond.L.Unlock()

	for b.pendingTxs.Len() == 0 {
		if err := b.pendingTxsCond.Wait(ctx); err != nil {
			return vmcore.Message{}, err
		}
	}

	return vmcore.Message{Type: vmcore.PendingTxs}, nil
}

func (b *builder) BuildBlock(ctx context.Context, blockContext *runtime.Runtime) (chain.Block, error) {
	preferredBlk, err := b.chain.GetBlock(b.preference)
	if err != nil {
		return nil, err
	}

	preferredState, err := preferredBlk.State()
	if err != nil {
		return nil, err
	}

	parentTimestamp := preferredBlk.Timestamp()
	timestamp := time.Now().Truncate(time.Second)
	if timestamp.Before(parentTimestamp) {
		timestamp = parentTimestamp
	}

	wipBlock := xsblock.Stateless{
		ParentID:  b.preference,
		Timestamp: timestamp.Unix(),
		Height:    preferredBlk.Height() + 1,
	}

	b.pendingTxsCond.L.Lock()
	defer b.pendingTxsCond.L.Unlock()

	currentState := versiondb.New(preferredState)
	for len(wipBlock.Txs) < MaxTxsPerBlock {
		txID, currentTx, exists := b.pendingTxs.Oldest()
		if !exists {
			break
		}
		b.pendingTxs.Delete(txID)

		sender, err := currentTx.SenderID()
		if err != nil {
			// This tx was invalid, drop it and continue block building
			continue
		}

		txState := versiondb.New(currentState)
		txExecutor := execute.Tx{
			Context:      ctx,
			Runtime:      b.chainRuntime,
			Database:     txState,
			BlockContext: blockContext,
			TxID:         txID,
			Sender:       sender,
		}
		if err := currentTx.Unsigned.Visit(&txExecutor); err != nil {
			// This tx was invalid, drop it and continue block building
			continue
		}
		if err := txState.Commit(); err != nil {
			return nil, err
		}

		wipBlock.Txs = append(wipBlock.Txs, currentTx)
	}
	return b.chain.NewBlock(&wipBlock)
}
