package actionsrunner

import (
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"math/big"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/service/automation"
	automationstore "github.com/tidbcloud/AgentTeam/internal/store/automation"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestAccessTokenExpiresBeforeSessionHeartbeat(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	control := &fakeAutomationControlPlane{}
	service := NewService(automationstore.New(nil), control, func() time.Time { return now })
	token, err := encodeAccessToken(accessTokenClaims{
		WorkspaceID: "ws-1",
		RunnerID:    "runner-1",
		SessionID:   "session-1",
		Scope:       accessTokenScopeBroker,
		IssuedAt:    now.Add(-2 * time.Hour).Unix(),
		ExpiresAt:   now.Add(-time.Hour).Unix(),
	})
	require.NoError(t, err)

	_, err = service.CreateBrokerSession(context.Background(), token, "runner-1")
	require.ErrorIs(t, err, automation.ErrUnauthorized)
	require.Zero(t, control.heartbeatCalls)
}

func TestAccessTokenAuthenticatesFromStoreAcrossServiceInstances(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	store := newActionsRunnerStoreForTest(t)
	token := createStoredAccessTokenForTest(t, store, now)
	require.Len(t, strings.Split(token, "."), 3)
	require.NotContains(t, token, "mxar_")

	control := &fakeAutomationControlPlane{}
	service := NewService(store, control, func() time.Time { return now })
	session, err := service.CreateBrokerSession(context.Background(), token, "")
	require.NoError(t, err)
	require.Equal(t, "session-1", session.SessionID)
	require.Equal(t, "runner-1", session.OwnerName)
	require.Equal(t, 1, control.heartbeatCalls)
	require.Equal(t, nativeSessionTokenHash("session-token-1"), control.heartbeatTokenHash)
}

func TestNextMessageHeartbeatsOnceBeforeAcquire(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	store := newActionsRunnerStoreForTest(t)
	token := createStoredAccessTokenForTest(t, store, now)
	control := &fakeAutomationControlPlane{}
	service := NewService(store, control, func() time.Time { return now })

	message, err := service.NextMessage(context.Background(), token, "http://meshx.test/_runs/")
	require.NoError(t, err)
	require.Nil(t, message)
	require.Equal(t, 1, control.heartbeatCalls)
	require.Equal(t, 1, control.acquireCalls)
}

func TestAccessTokenScopeEnforcedBeforeHeartbeat(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	store := newActionsRunnerStoreForTest(t)
	brokerToken := createStoredAccessTokenForTest(t, store, now)
	runServiceToken := createStoredScopedAccessTokenForTest(t, store, now, accessTokenScopeRunService, runServiceAccessTokenTTL)
	control := &fakeAutomationControlPlane{}
	service := NewService(store, control, func() time.Time { return now })

	_, err := service.CreateBrokerSession(context.Background(), runServiceToken, "")
	require.ErrorIs(t, err, automation.ErrUnauthorized)
	_, err = service.RenewJob(context.Background(), brokerToken, RenewJobRequest{JobID: queueJobGUID(1)})
	require.ErrorIs(t, err, automation.ErrUnauthorized)
	require.Zero(t, control.heartbeatCalls)
}

func TestVerifyClientAssertionRequiresBoundedExpiry(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	const clientID = "client-1"
	const audience = "http://meshx.test/_oauth/token"
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	publicKeyXML := clientAssertionPublicKeyXMLForTest(t, key)

	tests := []struct {
		name    string
		payload map[string]any
		wantErr bool
	}{
		{
			name: "valid",
			payload: map[string]any{
				"sub": clientID,
				"iss": clientID,
				"aud": audience,
				"nbf": now.Add(-time.Minute).Unix(),
				"exp": now.Add(time.Minute).Unix(),
			},
		},
		{
			name: "missing exp",
			payload: map[string]any{
				"sub": clientID,
				"iss": clientID,
				"aud": audience,
				"nbf": now.Add(-time.Minute).Unix(),
			},
			wantErr: true,
		},
		{
			name: "expired exp",
			payload: map[string]any{
				"sub": clientID,
				"iss": clientID,
				"aud": audience,
				"nbf": now.Add(-10 * time.Minute).Unix(),
				"exp": now.Add(-2 * time.Minute).Unix(),
			},
			wantErr: true,
		},
		{
			name: "overlong exp",
			payload: map[string]any{
				"sub": clientID,
				"iss": clientID,
				"aud": audience,
				"nbf": now.Add(-time.Minute).Unix(),
				"exp": now.Add(clientAssertionMaxTTL + 2*time.Minute).Unix(),
			},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assertion := signClientAssertionForTest(t, key, tt.payload)
			claims, err := parseClientAssertion(assertion)
			require.NoError(t, err)
			err = verifyClientAssertion(assertion, claims, clientID, publicKeyXML, audience, now)
			if tt.wantErr {
				require.ErrorIs(t, err, automation.ErrUnauthorized)
				return
			}
			require.NoError(t, err)
		})
	}
}

func TestAcquireJobDoesNotPersistRunServiceTokenBeforeAccept(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	store, db := newActionsRunnerStoreAndDBForTest(t)
	brokerToken := createStoredAccessTokenForTest(t, store, now)
	control := &fakeAutomationControlPlane{acceptErr: automation.ErrNotFound}
	service := NewService(store, control, func() time.Time { return now })

	_, err := service.AcquireJob(context.Background(), brokerToken, AcquireJobRequest{JobMessageID: "123", RunServiceURL: "http://meshx.test/_runs/"})
	require.ErrorIs(t, err, automation.ErrNotFound)

	var tokenCount int64
	require.NoError(t, db.Model(&automationstore.AutomationRunnerToken{}).Where("token_kind = ?", automationstore.AutomationRunnerTokenKindActionsAccess).Count(&tokenCount).Error)
	require.Equal(t, int64(1), tokenCount)
}

func createStoredAccessTokenForTest(t *testing.T, store *automationstore.Store, now time.Time) string {
	t.Helper()
	return createStoredScopedAccessTokenForTest(t, store, now, accessTokenScopeBroker, brokerAccessTokenTTL)
}

func createStoredScopedAccessTokenForTest(t *testing.T, store *automationstore.Store, now time.Time, scope string, ttl time.Duration) string {
	t.Helper()
	token, err := encodeAccessToken(accessTokenClaims{
		WorkspaceID: "ws-1",
		RunnerID:    "runner-1",
		SessionID:   "session-1",
		Scope:       scope,
		IssuedAt:    now.Unix(),
		ExpiresAt:   now.Add(ttl).Unix(),
	})
	require.NoError(t, err)
	require.NotContains(t, token, "session-token-1")
	require.NotContains(t, tokenPayload(t, token), "session_token")
	require.NotContains(t, tokenPayload(t, token), "session-token-1")
	require.NoError(t, store.CreateActionsRunnerAccessToken(context.Background(), automationstore.AutomationRunnerToken{
		TokenHash:        accessTokenHash(token),
		WorkspaceID:      "ws-1",
		RunnerID:         "runner-1",
		TokenScope:       scope,
		SessionID:        "session-1",
		SessionTokenHash: nativeSessionTokenHash("session-token-1"),
		ExpiresAt:        now.Add(ttl),
		CreatedAt:        now,
		UpdatedAt:        now,
	}))
	return token
}

func TestRunServiceTokenOutlivesBrokerTokenBudget(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 0, 0, 0, time.UTC)
	store := newActionsRunnerStoreForTest(t)
	service := NewService(store, &fakeAutomationControlPlane{}, func() time.Time { return now })

	token, err := service.issueAccessToken(context.Background(), "ws-1", "runner-1", "session-1", nativeSessionTokenHash("session-token-1"), accessTokenScopeRunService, runServiceAccessTokenTTL)
	require.NoError(t, err)
	require.NotContains(t, tokenPayload(t, token), "session_token")
	require.NotContains(t, tokenPayload(t, token), "session-token-1")

	afterBrokerTTL := now.Add(brokerAccessTokenTTL + time.Hour)
	claims, err := decodeAccessToken(token, afterBrokerTTL)
	require.NoError(t, err)
	require.Equal(t, "session-1", claims.SessionID)
	require.Equal(t, accessTokenScopeRunService, claims.Scope)
	_, err = store.GetActiveActionsRunnerAccessTokenByHash(context.Background(), accessTokenHash(token), afterBrokerTTL)
	require.NoError(t, err)
}

func TestSuccessfulConclusionRequiresExactZero(t *testing.T) {
	tests := []struct {
		name string
		in   any
		want bool
	}{
		{name: "numeric zero", in: float64(0), want: true},
		{name: "numeric fractional", in: float64(0.5), want: false},
		{name: "numeric failure", in: float64(2), want: false},
		{name: "string success", in: "success", want: true},
		{name: "string zero", in: "0", want: true},
		{name: "string failure", in: "failed", want: false},
		{name: "missing", in: nil, want: false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.want, successfulConclusion(tt.in))
		})
	}
}

func TestAgentJobRequestMessageRejectsEnvForOfficialRunner(t *testing.T) {
	now := time.Date(2026, 6, 6, 10, 30, 0, 0, time.UTC)
	run := automationstore.AutomationRun{ID: "wfr-actions-env", WorkspaceID: "ws-1", RepositoryFullName: "pingcap/agentteam", DefinitionPath: ".agentteam/deploy.yml", GitRef: "main", CommitSHA: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Name: "Deploy"}
	job := automationstore.AutomationJob{ID: 42, Attempt: 1}

	cases := []struct {
		name      string
		inputJSON string
		want      string
	}{
		{
			name:      "job env",
			inputJSON: `{"env":{"TOKEN":"${{ secrets.deploy_token }}"},"steps":[{"run":"echo deploy"}]}`,
			want:      "automation job env is not supported",
		},
		{
			name:      "step env",
			inputJSON: `{"steps":[{"run":"echo deploy","env":{"TOKEN":"${{ secrets.deploy_token }}"}}]}`,
			want:      "automation job deploy step 0 env is not supported",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			jobRun := automationstore.AutomationJobRun{ID: "wfj-actions-env", JobID: "deploy", InputJSON: tc.inputJSON}
			_, err := agentJobRequestMessage("http://meshx.test/_runs/", now, run, job, jobRun)
			require.Error(t, err)
			require.Contains(t, err.Error(), tc.want)
		})
	}
}

func newActionsRunnerStoreForTest(t *testing.T) *automationstore.Store {
	t.Helper()
	store, _ := newActionsRunnerStoreAndDBForTest(t)
	return store
}

func newActionsRunnerStoreAndDBForTest(t *testing.T) (*automationstore.Store, *gorm.DB) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&automationstore.AutomationRunnerToken{}))
	return automationstore.New(db), db
}

func tokenPayload(t *testing.T, token string) string {
	t.Helper()
	parts := strings.Split(strings.TrimSpace(token), ".")
	require.Len(t, parts, 3)
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	require.NoError(t, err)
	return string(payload)
}

func clientAssertionPublicKeyXMLForTest(t *testing.T, key *rsa.PrivateKey) string {
	t.Helper()
	type rsaKeyValue struct {
		XMLName  xml.Name `xml:"RSAKeyValue"`
		Modulus  string   `xml:"Modulus"`
		Exponent string   `xml:"Exponent"`
	}
	encoded, err := xml.Marshal(rsaKeyValue{
		Modulus:  base64.StdEncoding.EncodeToString(key.PublicKey.N.Bytes()),
		Exponent: base64.StdEncoding.EncodeToString(big.NewInt(int64(key.PublicKey.E)).Bytes()),
	})
	require.NoError(t, err)
	return string(encoded)
}

func signClientAssertionForTest(t *testing.T, key *rsa.PrivateKey, payload map[string]any) string {
	t.Helper()
	header := base64.RawURLEncoding.EncodeToString([]byte(`{"typ":"JWT","alg":"PS256"}`))
	body, err := json.Marshal(payload)
	require.NoError(t, err)
	signingInput := header + "." + base64.RawURLEncoding.EncodeToString(body)
	sum := sha256.Sum256([]byte(signingInput))
	signature, err := rsa.SignPSS(rand.Reader, key, crypto.SHA256, sum[:], &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash})
	require.NoError(t, err)
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature)
}

type fakeAutomationControlPlane struct {
	heartbeatCalls     int
	heartbeatTokenHash string
	acquireCalls       int
	acceptErr          error
}

func (f *fakeAutomationControlPlane) LookupAutomationRunnerRegistrationToken(context.Context, string) (automation.AutomationRunnerRegistrationLookup, error) {
	return automation.AutomationRunnerRegistrationLookup{}, nil
}

func (f *fakeAutomationControlPlane) RegisterAutomationActionsRunnerCredential(context.Context, string, automationstore.AutomationActionsRunnerCredential) error {
	return nil
}

func (f *fakeAutomationControlPlane) CreateOrReuseAutomationRunnerSessionForRunner(context.Context, string, automation.AutomationRunnerSessionMutation) (automation.AutomationRunnerSession, error) {
	return automation.AutomationRunnerSession{}, nil
}

func (f *fakeAutomationControlPlane) HeartbeatAutomationRunnerSessionByTokenHash(_ context.Context, _ string, tokenHash string, _ automation.AutomationRunnerHeartbeatMutation) (automation.AutomationRunnerHeartbeat, error) {
	f.heartbeatCalls++
	f.heartbeatTokenHash = tokenHash
	return automation.AutomationRunnerHeartbeat{}, nil
}

func (f *fakeAutomationControlPlane) RevokeAutomationRunnerSessionByTokenHash(context.Context, string, string, automation.AutomationRunnerSessionMutation) error {
	return nil
}

func (f *fakeAutomationControlPlane) AcquireAutomationRunnerJobsByTokenHash(context.Context, string, string, automation.AutomationRunnerAcquireMutation) ([]automation.AutomationRunnerJob, error) {
	f.acquireCalls++
	return nil, nil
}

func (f *fakeAutomationControlPlane) AcceptAutomationRunnerJobByTokenHash(context.Context, string, string, automation.AutomationRunnerJobMutation) (automationstore.AutomationJobRun, error) {
	if f.acceptErr != nil {
		return automationstore.AutomationJobRun{}, f.acceptErr
	}
	return automationstore.AutomationJobRun{}, nil
}

func (f *fakeAutomationControlPlane) RenewAutomationRunnerJobByTokenHash(context.Context, string, string, automation.AutomationRunnerJobRenewMutation) (automation.AutomationRunnerJobLease, error) {
	return automation.AutomationRunnerJobLease{}, nil
}

func (f *fakeAutomationControlPlane) CompleteAutomationRunnerJobByTokenHash(context.Context, string, string, automation.AutomationRunnerJobMutation) (automationstore.AutomationJobRun, error) {
	return automationstore.AutomationJobRun{}, nil
}

func (f *fakeAutomationControlPlane) FailAutomationRunnerJobByTokenHash(context.Context, string, string, automation.AutomationRunnerJobMutation) (automationstore.AutomationJobRun, error) {
	return automationstore.AutomationJobRun{}, nil
}