package miner

import (
	"encoding/hex"
	"errors"
	"fmt"
	"sync"

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

var (
	ErrEngineNotReady = errors.New("randomx VM not initialized")
	ErrBlobTooShort   = errors.New("blob shorter than nonce offset")
)

const nonceOffset = 39
const nonceSize = 4

// go-randomx is a pure-Go implementation; hardware flags are ignored internally.
const randomxFlags = 0

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}
}

// Reset clears the VM so the next SetJob reinitializes RandomX state.
func (e *Engine) Reset() {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.vm = nil
	e.seedHex = ""
	e.blob = nil
}

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)
		// go-randomx requires SuperScalar programs to be built separately after
		// seeding the cache; Randomx_init_cache only populates the Argon2d blocks.
		// Without this step every CalculateHash call crashes with a nil-pointer.
		gen := randomx.Init_Blake2Generator(seed, 0)
		for i := range e.cache.Programs {
			e.cache.Programs[i] = randomx.Build_SuperScalar_Program(gen)
		}
		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 {
		return "", "", ErrEngineNotReady
	}
	if len(e.blob) < nonceOffset+nonceSize {
		return "", "", fmt.Errorf("%w (need %d bytes, have %d)", ErrBlobTooShort, nonceOffset+nonceSize, len(e.blob))
	}

	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
}