package auth

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"strings"
	"testing"
	"time"
)

func TestGenerateAndValidateJobToken_RoundTrip(t *testing.T) {
	token, err := GenerateJobToken(10, 20, 30)
	if err != nil {
		t.Fatalf("GenerateJobToken failed: %v", err)
	}
	claims, err := ValidateJobToken(token)
	if err != nil {
		t.Fatalf("ValidateJobToken failed: %v", err)
	}
	if claims.JobID != 10 || claims.RunnerID != 20 || claims.TaskID != 30 {
		t.Fatalf("unexpected claims: %+v", claims)
	}
}

func TestValidateJobToken_RejectsTampering(t *testing.T) {
	token, err := GenerateJobToken(1, 2, 3)
	if err != nil {
		t.Fatalf("GenerateJobToken failed: %v", err)
	}
	parts := strings.Split(token, ".")
	if len(parts) != 2 {
		t.Fatalf("unexpected token format: %q", token)
	}

	rawClaims, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		t.Fatalf("decode claims failed: %v", err)
	}
	var claims JobTokenClaims
	if err := json.Unmarshal(rawClaims, &claims); err != nil {
		t.Fatalf("unmarshal claims failed: %v", err)
	}
	claims.JobID = 999
	tamperedClaims, _ := json.Marshal(claims)
	tampered := base64.RawURLEncoding.EncodeToString(tamperedClaims) + "." + parts[1]

	if _, err := ValidateJobToken(tampered); err == nil {
		t.Fatal("expected signature validation error for tampered token")
	}
}

func TestValidateJobToken_RejectsExpired(t *testing.T) {
	expiredClaims := JobTokenClaims{
		JobID:    1,
		RunnerID: 2,
		TaskID:   3,
		Exp:      time.Now().Add(-time.Minute).Unix(),
	}
	claimsJSON, _ := json.Marshal(expiredClaims)
	sigToken, err := GenerateJobToken(1, 2, 3)
	if err != nil {
		t.Fatalf("GenerateJobToken failed: %v", err)
	}
	parts := strings.Split(sigToken, ".")
	if len(parts) != 2 {
		t.Fatalf("unexpected token format: %q", sigToken)
	}
	// Re-sign expired payload with package secret.
	h := signClaimsForTest(claimsJSON)
	expiredToken := base64.RawURLEncoding.EncodeToString(claimsJSON) + "." + base64.RawURLEncoding.EncodeToString(h)

	if _, err := ValidateJobToken(expiredToken); err == nil {
		t.Fatal("expected token expiration error")
	}
}

func signClaimsForTest(claims []byte) []byte {
	h := hmac.New(sha256.New, jobTokenSecret)
	_, _ = h.Write(claims)
	return h.Sum(nil)
}