package miner

import (
	"encoding/hex"
	"sync"

	"git.gammaspectra.live/P2Pool/go-randomx"
)

const nonceOffset = 39
const nonceSize = 4

// RandomX JIT + hardware AES for best hashrate on supported CPUs.
const randomxFlags = 10 // RANDOMX_FLAG_HARD_AES (2) | RANDOMX_FLAG_JIT (8)

type Engine struct {
	mu        sync.RWMutex
	cache     *randomx.Randomx_Cache
	vm        *randomx.VM
	seedHex   string
	blob      []byte
}

func NewEngine() *Engine {
	cache := randomx.Randomx_alloc_cache(randomxFlags)
	return &Engine{cache: cache}
}

func (e *Engine) SetJob(seedHex, blobHex string) error {
	seed, err := hex.DecodeString(seedHex)
	if err != nil {
		return err
	}
	blob, err := hex.DecodeString(blobHex)
	if err != nil {
		return err
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	if e.seedHex != seedHex {
		e.cache.Randomx_init_cache(seed)
		e.vm = e.cache.VM_Initialize()
		e.seedHex = seedHex
	}
	e.blob = append([]byte(nil), blob...)
	return nil
}

func (e *Engine) HashAtNonce(nonce uint32) (hashHex string, blobHex string, err error) {
	e.mu.RLock()
	defer e.mu.RUnlock()

	if e.vm == nil || len(e.blob) < nonceOffset+nonceSize {
		return "", "", nil
	}

	work := append([]byte(nil), e.blob...)
	work[nonceOffset] = byte(nonce)
	work[nonceOffset+1] = byte(nonce >> 8)
	work[nonceOffset+2] = byte(nonce >> 16)
	work[nonceOffset+3] = byte(nonce >> 24)

	out := make([]byte, 32)
	e.vm.CalculateHash(work, out)
	return hex.EncodeToString(out), hex.EncodeToString(work), nil
}