package distworker

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"rocketgit.ru/rsmon/worker/internal/wire"
)

func signedInit(t *testing.T, private ed25519.PrivateKey, account, version int64) *wire.WorkerInit {
	t.Helper()
	init := &wire.WorkerInit{WorkerID: "w", SigningKeyID: "v1", AccountID: &account, ConfigVersion: version, ExpiresAt: time.Now().Add(time.Minute).Format(time.RFC3339Nano)}
	body, err := json.Marshal(init)
	require.NoError(t, err)
	init.Signature = base64.StdEncoding.EncodeToString(ed25519.Sign(private, body))
	return init
}

func TestVerifyInitRejectsTamperWrongKeyAndAccountChange(t *testing.T) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	r := NewRunner(&Config{})
	r.state = workerState{WorkerID: "w", VerificationKey: base64.StdEncoding.EncodeToString(pub), SigningKeyID: "v1", AccountID: 1}
	init := signedInit(t, priv, 1, 1)
	require.NoError(t, r.verifyInit(init))
	init.ConfigVersion = 2
	assert.Error(t, r.verifyInit(init), "tamper must invalidate signature")
	_, other, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	assert.Error(t, r.verifyInit(signedInit(t, other, 1, 2)), "wrong key")
	assert.Error(t, r.verifyInit(signedInit(t, priv, 2, 2)), "account scope must be immutable")
	unknown := signedInit(t, priv, 1, 2)
	unknown.SigningKeyID = "unknown"
	assert.Error(t, r.verifyInit(unknown), "unknown key id")
}

func TestVerifyInitAcceptsIdenticalRestartReplayOnly(t *testing.T) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	first := NewRunner(&Config{})
	first.state = workerState{WorkerID: "w", VerificationKey: base64.StdEncoding.EncodeToString(pub), SigningKeyID: "v1", AccountID: 1}
	init := signedInit(t, priv, 1, 4)
	require.NoError(t, first.verifyInit(init))
	// Restart restores only durable identity state. The exact verified snapshot
	// remains safe to replay at the same generation.
	restarted := NewRunner(&Config{})
	restarted.state = first.state
	require.NoError(t, restarted.verifyInit(signedInit(t, priv, 1, 4)))
	assert.Error(t, restarted.verifyInit(signedInit(t, priv, 1, 3)), "older generation")
	different := signedInit(t, priv, 1, 4)
	different.RegionCode = "other"
	body, err := json.Marshal(&wire.WorkerInit{WorkerID: different.WorkerID, SigningKeyID: different.SigningKeyID, AccountID: different.AccountID, ConfigVersion: different.ConfigVersion, ExpiresAt: different.ExpiresAt, RegionCode: different.RegionCode})
	require.NoError(t, err)
	different.Signature = base64.StdEncoding.EncodeToString(ed25519.Sign(priv, body))
	assert.Error(t, restarted.verifyInit(different), "same generation different content")
	require.NoError(t, restarted.verifyInit(signedInit(t, priv, 1, 5)), "newer generation")
}