package token

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestExecTokenSignedClaimsAndActiveRecord(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_800_000_000, 0)
	manager.now = func() time.Time { return now }
	value, issued, err := manager.GenerateExecToken(ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-1",
		ProcessGeneration: "generation-1",
		TTL:               time.Hour,
	})
	require.NoError(t, err)

	parts := strings.Split(value, ".")
	require.Len(t, parts, 3)
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	require.NoError(t, err)
	var claims map[string]any
	require.NoError(t, json.Unmarshal(payload, &claims))
	require.ElementsMatch(t, []string{"version", "agent_id", "jti", "exp"}, mapKeys(claims))
	require.NotContains(t, claims, "server_id")
	require.NotContains(t, claims, "session_id")
	require.NotContains(t, claims, "run_id")
	require.NotContains(t, claims, "process_generation")
	require.NotContains(t, claims, "extension_catalog_version")
	require.NotContains(t, claims, "extension_catalog_digest")
	require.Equal(t, issued.JTI, claims["jti"])

	validated, err := manager.ValidateExecToken("agent-1", "server-1", value)
	require.NoError(t, err)
	require.Equal(t, "session-1", validated.SessionID)
	require.Equal(t, "run-1", validated.RunID)
	require.Equal(t, "generation-1", validated.ProcessGeneration)
}

func TestExecTokenRejectsTamperMismatchExpiryAndRevoke(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_800_000_000, 0)
	manager.now = func() time.Time { return now }
	value, _, err := manager.GenerateExecToken(ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-1",
		TTL: time.Minute,
	})
	require.NoError(t, err)

	parts := strings.Split(value, ".")
	signature, err := base64.RawURLEncoding.DecodeString(parts[2])
	require.NoError(t, err)
	signature[0] ^= 0xff
	parts[2] = base64.RawURLEncoding.EncodeToString(signature)
	tampered := strings.Join(parts, ".")
	_, err = manager.ValidateExecToken("agent-1", "server-1", tampered)
	require.Error(t, err)
	_, err = manager.ValidateExecToken("agent-2", "server-1", value)
	require.ErrorIs(t, err, ErrAgentMismatch)
	_, err = manager.ValidateExecToken("agent-1", "server-2", value)
	require.ErrorIs(t, err, ErrAgentTokenServerMismatch)

	now = now.Add(2 * time.Minute)
	_, err = manager.ValidateExecToken("agent-1", "server-1", value)
	require.ErrorIs(t, err, ErrExpired)

	now = time.Unix(1_800_000_000, 0)
	value, issued, err := manager.GenerateExecToken(ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-2",
	})
	require.NoError(t, err)
	require.True(t, manager.RevokeExecToken(issued.JTI))
	_, err = manager.ValidateExecToken("agent-1", "server-1", value)
	require.True(t, errors.Is(err, ErrAgentTokenRevoked) || errors.Is(err, ErrAgentTokenInactive))
}

func TestExecTokenRejectsNonMinimalBase64URL(t *testing.T) {
	manager := newTestManager(t)
	value, _, err := manager.GenerateExecToken(ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-1",
	})
	require.NoError(t, err)

	parts := strings.Split(value, ".")
	expectedSignature, err := base64.RawURLEncoding.DecodeString(parts[2])
	require.NoError(t, err)
	const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
	last := strings.IndexByte(alphabet, parts[2][len(parts[2])-1])
	require.GreaterOrEqual(t, last, 0)
	require.Equal(t, 0, last&0x0f, "64-byte Ed25519 signature has four unused trailing base64 bits")
	parts[2] = parts[2][:len(parts[2])-1] + string(alphabet[last+1])
	nonMinimal := strings.Join(parts, ".")
	decodedSignature, err := base64.RawURLEncoding.DecodeString(parts[2])
	require.NoError(t, err)
	require.Equal(t, expectedSignature, decodedSignature, "test must alter only unused base64 bits")

	_, err = manager.ValidateExecToken("agent-1", "server-1", nonMinimal)
	require.ErrorIs(t, err, ErrInvalidFormat)
}

func TestExecTokenExpiryUsesOneSecondGranularity(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_800_000_000, 900_000_000)
	manager.now = func() time.Time { return now }
	value, issued, err := manager.GenerateExecToken(ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-1",
		TTL: 2 * time.Second,
	})
	require.NoError(t, err)
	require.Equal(t, time.Unix(1_800_000_002, 0), issued.ExpiresAt)

	now = issued.ExpiresAt.Add(-time.Nanosecond)
	validated, err := manager.ValidateExecToken("agent-1", "server-1", value)
	require.NoError(t, err)
	require.Equal(t, issued.ExpiresAt, validated.ExpiresAt)

	now = issued.ExpiresAt
	_, err = manager.ValidateExecToken("agent-1", "server-1", value)
	require.ErrorIs(t, err, ErrExpired)
}

func TestExecTokenRejectsTTLThatCannotReachSignedSecond(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_800_000_000, 100_000_000)
	manager.now = func() time.Time { return now }

	_, _, err := manager.GenerateExecToken(ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-1",
		TTL: 100 * time.Millisecond,
	})
	require.ErrorContains(t, err, "next signed second")
	require.Empty(t, manager.execTokens)
}

func TestExecTokenCapacityFailsWithoutEvictingActiveToken(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_800_000_000, 0)
	manager.now = func() time.Time { return now }
	request := ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1",
		TTL: time.Hour,
	}

	var firstValue string
	var firstInfo ExecTokenInfo
	for index := 0; index < maxExecTokens; index++ {
		request.RunID = "run-" + fmt.Sprint(index)
		value, info, err := manager.GenerateExecToken(request)
		require.NoError(t, err)
		if index == 0 {
			firstValue, firstInfo = value, info
		}
	}
	request.RunID = "over-capacity"
	_, _, err := manager.GenerateExecToken(request)
	require.ErrorIs(t, err, ErrExecTokenCapacity)
	require.Len(t, manager.execTokens, maxExecTokens)

	_, err = manager.ValidateExecToken("agent-1", "server-1", firstValue)
	require.NoError(t, err, "capacity pressure must not silently evict an active runtime")

	require.True(t, manager.RevokeExecToken(firstInfo.JTI))
	value, _, err := manager.GenerateExecToken(request)
	require.NoError(t, err, "revoked records must be pruned before capacity is enforced")
	require.Len(t, manager.execTokens, maxExecTokens)
	_, err = manager.ValidateExecToken("agent-1", "server-1", value)
	require.NoError(t, err)
}

func TestExecTokenGenerationPrunesExpiredAndRevokedRecords(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_800_000_000, 0)
	manager.now = func() time.Time { return now }
	request := ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1",
		RunID: "run-1", TTL: 2 * time.Second,
	}
	expiredValue, expiredInfo, err := manager.GenerateExecToken(request)
	require.NoError(t, err)

	request.RunID = "run-2"
	_, revokedInfo, err := manager.GenerateExecToken(request)
	require.NoError(t, err)
	require.True(t, manager.RevokeExecToken(revokedInfo.JTI))
	require.Contains(t, manager.execTokens, expiredInfo.JTI)
	require.Contains(t, manager.execTokens, revokedInfo.JTI)

	now = now.Add(3 * time.Second)
	request.RunID = "run-3"
	request.TTL = time.Hour
	_, activeInfo, err := manager.GenerateExecToken(request)
	require.NoError(t, err)
	require.Equal(t, map[string]execTokenRecord{
		activeInfo.JTI: manager.execTokens[activeInfo.JTI],
	}, manager.execTokens)
	_, err = manager.ValidateExecToken("agent-1", "server-1", expiredValue)
	require.ErrorIs(t, err, ErrExpired)
}

func TestExecTokenRotatedGenerationSurvivesOriginalTwelveHourExpiry(t *testing.T) {
	manager := newTestManager(t)
	now := time.Unix(1_900_000_000, 0)
	manager.now = func() time.Time { return now }
	request := ExecTokenRequest{
		AgentID: "agent-1", ServerID: "server-1", SessionID: "session-1", RunID: "run-1",
		TTL: 12 * time.Hour,
	}
	initial, _, err := manager.GenerateExecToken(request)
	require.NoError(t, err)

	now = now.Add(6 * time.Hour)
	rotated, _, err := manager.GenerateExecToken(request)
	require.NoError(t, err)

	now = now.Add(7 * time.Hour)
	_, err = manager.ValidateExecToken("agent-1", "server-1", initial)
	require.ErrorIs(t, err, ErrExpired)
	_, err = manager.ValidateExecToken("agent-1", "server-1", rotated)
	require.NoError(t, err)
}

func mapKeys(input map[string]any) []string {
	keys := make([]string, 0, len(input))
	for key := range input {
		keys = append(keys, key)
	}
	return keys
}