package multithreaded

import (
	"encoding/binary"
	"io"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"

	"github.com/ethereum-optimism/optimism/cannon/mipsevm"
	"github.com/ethereum-optimism/optimism/cannon/mipsevm/arch"
	"github.com/ethereum-optimism/optimism/cannon/mipsevm/exec"
)

const (
	THREAD_ID_STATE_WITNESS_OFFSET           = 0
	THREAD_EXIT_CODE_WITNESS_OFFSET          = THREAD_ID_STATE_WITNESS_OFFSET + arch.WordSizeBytes
	THREAD_EXITED_WITNESS_OFFSET             = THREAD_EXIT_CODE_WITNESS_OFFSET + 1
	THREAD_FUTEX_ADDR_WITNESS_OFFSET         = THREAD_EXITED_WITNESS_OFFSET + 1
	THREAD_FUTEX_VAL_WITNESS_OFFSET          = THREAD_FUTEX_ADDR_WITNESS_OFFSET + arch.WordSizeBytes
	THREAD_FUTEX_TIMEOUT_STEP_WITNESS_OFFSET = THREAD_FUTEX_VAL_WITNESS_OFFSET + arch.WordSizeBytes
	THREAD_FUTEX_CPU_WITNESS_OFFSET          = THREAD_FUTEX_TIMEOUT_STEP_WITNESS_OFFSET + 8
	THREAD_REGISTERS_WITNESS_OFFSET          = THREAD_FUTEX_CPU_WITNESS_OFFSET + (4 * arch.WordSizeBytes)

	// SERIALIZED_THREAD_SIZE is the size of a serialized ThreadState object
	// 166 and 322 bytes for 32 and 64-bit respectively
	SERIALIZED_THREAD_SIZE = THREAD_REGISTERS_WITNESS_OFFSET + (32 * arch.WordSizeBytes)

	// THREAD_WITNESS_SIZE is the size of a thread witness encoded in bytes.
	//
	//	It consists of the active thread serialized and concatenated with the
	//	32 byte hash onion of the active thread stack without the active thread
	THREAD_WITNESS_SIZE = SERIALIZED_THREAD_SIZE + 32
)

// The empty thread root - keccak256(bytes32(0) ++ bytes32(0))
var EmptyThreadsRoot common.Hash = common.HexToHash("0xad3228b676f7d3cd4284a5443f17f1962b36e491b30a40b2405849e597ba5fb5")

type ThreadState struct {
	ThreadId         Word               `json:"threadId"`
	ExitCode         uint8              `json:"exit"`
	Exited           bool               `json:"exited"`
	FutexAddr        Word               `json:"futexAddr"`
	FutexVal         Word               `json:"futexVal"`
	FutexTimeoutStep uint64             `json:"futexTimeoutStep"`
	Cpu              mipsevm.CpuScalars `json:"cpu"`
	Registers        [32]Word           `json:"registers"`
}

func CreateEmptyThread() *ThreadState {
	initThreadId := Word(0)
	return &ThreadState{
		ThreadId: initThreadId,
		ExitCode: 0,
		Exited:   false,
		Cpu: mipsevm.CpuScalars{
			PC:     0,
			NextPC: 4,
			LO:     0,
			HI:     0,
		},
		FutexAddr:        exec.FutexEmptyAddr,
		FutexVal:         0,
		FutexTimeoutStep: 0,
		Registers:        [32]Word{},
	}
}

func (t *ThreadState) serializeThread() []byte {
	out := make([]byte, 0, SERIALIZED_THREAD_SIZE)

	out = arch.ByteOrderWord.AppendWord(out, t.ThreadId)
	out = append(out, t.ExitCode)
	out = mipsevm.AppendBoolToWitness(out, t.Exited)
	out = arch.ByteOrderWord.AppendWord(out, t.FutexAddr)
	out = arch.ByteOrderWord.AppendWord(out, t.FutexVal)
	out = binary.BigEndian.AppendUint64(out, t.FutexTimeoutStep)

	out = arch.ByteOrderWord.AppendWord(out, t.Cpu.PC)
	out = arch.ByteOrderWord.AppendWord(out, t.Cpu.NextPC)
	out = arch.ByteOrderWord.AppendWord(out, t.Cpu.LO)
	out = arch.ByteOrderWord.AppendWord(out, t.Cpu.HI)

	for _, r := range t.Registers {
		out = arch.ByteOrderWord.AppendWord(out, r)
	}

	return out
}

// Serialize writes the ThreadState in a simple binary format which can be read again using Deserialize
// The format exactly matches the serialization generated by serializeThread used for thread proofs.
func (t *ThreadState) Serialize(out io.Writer) error {
	_, err := out.Write(t.serializeThread())
	return err
}

func (t *ThreadState) Deserialize(in io.Reader) error {
	if err := binary.Read(in, binary.BigEndian, &t.ThreadId); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.ExitCode); err != nil {
		return err
	}
	var exited uint8
	if err := binary.Read(in, binary.BigEndian, &exited); err != nil {
		return err
	}
	t.Exited = exited != 0
	if err := binary.Read(in, binary.BigEndian, &t.FutexAddr); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.FutexVal); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.FutexTimeoutStep); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.Cpu.PC); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.Cpu.NextPC); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.Cpu.LO); err != nil {
		return err
	}
	if err := binary.Read(in, binary.BigEndian, &t.Cpu.HI); err != nil {
		return err
	}
	// Read the registers as big endian Words
	for i := range t.Registers {
		if err := binary.Read(in, binary.BigEndian, &t.Registers[i]); err != nil {
			return err
		}
	}
	return nil
}

func computeThreadRoot(prevStackRoot common.Hash, threadToPush *ThreadState) common.Hash {
	hashedThread := crypto.Keccak256Hash(threadToPush.serializeThread())

	var hashData []byte
	hashData = append(hashData, prevStackRoot[:]...)
	hashData = append(hashData, hashedThread[:]...)

	return crypto.Keccak256Hash(hashData)
}
