package extensionflow

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/extension/binding"
	"github.com/tidbcloud/AgentTeam/internal/extension/egress"
	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/store/app"
	"github.com/tidbcloud/AgentTeam/internal/store/serviceactor"
)

type recordingInvocationEgress struct {
	mu          sync.Mutex
	calls       int
	bounds      []binding.BoundRequest
	credentials []egress.Credential
	results     []egress.Result
	errors      []error
}

func (executor *recordingInvocationEgress) Execute(
	_ context.Context,
	bound binding.BoundRequest,
	credential egress.Credential,
) (egress.Result, error) {
	executor.mu.Lock()
	defer executor.mu.Unlock()
	index := executor.calls
	executor.calls++
	executor.bounds = append(executor.bounds, bound)
	executor.credentials = append(executor.credentials, credential)
	if index < len(executor.results) {
		var err error
		if index < len(executor.errors) {
			err = executor.errors[index]
		}
		return executor.results[index], err
	}
	return egress.Result{}, errors.New("unexpected egress call")
}

func (executor *recordingInvocationEgress) callCount() int {
	executor.mu.Lock()
	defer executor.mu.Unlock()
	return executor.calls
}

type recordingOAuthRefresher struct {
	mu       sync.Mutex
	calls    int
	requests []OAuthRefreshRequest
	result   OAuthCredential
	err      error
}

func (refresher *recordingOAuthRefresher) RefreshExtensionOAuth(
	_ context.Context,
	request OAuthRefreshRequest,
) (OAuthCredential, error) {
	refresher.mu.Lock()
	defer refresher.mu.Unlock()
	refresher.calls++
	refresher.requests = append(refresher.requests, request)
	return cloneOAuthCredential(refresher.result), refresher.err
}

func TestInvocationAuditTimeoutDefaultsAndOverrideStayPerService(t *testing.T) {
	first := &Service{}
	second := &Service{}

	require.Equal(t, 5*time.Second, first.auditWriteTimeout())
	require.Equal(t, 5*time.Second, second.auditWriteTimeout())

	first.invocationAuditTimeout = 30 * time.Second
	require.Equal(t, 30*time.Second, first.auditWriteTimeout())
	require.Equal(t, 5*time.Second, second.auditWriteTimeout())
}

func TestInvokeAuthoritativeSuccessPersistsSecretFreeAudit(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_success")
	now := time.Date(2026, 8, 11, 9, 0, 0, 0, time.UTC)
	executor := &recordingInvocationEgress{
		results: []egress.Result{{
			Data:           json.RawMessage(`{"items":[1,2]}`),
			UpstreamStatus: 200,
			StatusClass:    "2xx",
			ContentType:    "application/json",
			ResponseBytes:  15,
			FinalBytes:     50,
		}},
		errors: []error{nil},
	}
	service := New(
		db,
		fixture.Registry,
		[]byte(testCredentialKey),
		Options{
			Now: func() time.Time { return now },
			EgressFactory: func(
				registry.ActionPlan,
			) (InvocationEgress, error) {
				return executor, nil
			},
		},
	)
	host, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"invoke-success",
		fixture.ToolsDigest,
		"tools-primary",
		"agent-1",
		map[string]any{
			"accessToken": "provider-secret",
			"scopes":      []string{},
		},
	)

	result, err := service.Invoke(
		ctx,
		invocationParams(
			fixture,
			"agent-1",
			"tools-primary",
			"list-items",
			fixture.ToolsDigest,
			version,
		),
	)
	require.NoError(t, err)
	require.Equal(t, "tools-primary", result.Domain)
	require.Equal(t, "list-items", result.Action)
	require.JSONEq(t, `{"items":[1,2]}`, string(result.Data))
	require.Equal(t, 1, executor.callCount())
	require.Equal(t, "provider-secret", executor.credentials[0].Value)
	require.Equal(
		t,
		"https://api.ext-test.example.com/items",
		executor.bounds[0].URL.String(),
	)

	var (
		count          int
		selectionID    string
		extensionName  string
		classification string
		requestHash    string
	)
	require.NoError(t, db.QueryRow(ctx, `
SELECT COUNT(*), MAX(service_installation_extension_id),
       MAX(extension_name), MAX(result_classification), MAX(request_hash)
FROM extension_invocation_audits
WHERE workspace_id=$1 AND agent_id=$2`,
		fixture.WorkspaceID,
		"agent-1",
	).Scan(
		&count,
		&selectionID,
		&extensionName,
		&classification,
		&requestHash,
	))
	require.Equal(t, 1, count)
	require.Equal(t, host.Selection.ID, selectionID)
	require.Equal(t, "tools-primary", extensionName)
	require.Equal(t, "success", classification)
	require.Regexp(t, `^sha256:[0-9a-f]{64}$`, requestHash)

	var secretBearingRows int
	require.NoError(t, db.QueryRow(ctx, `
SELECT COUNT(*) FROM extension_invocation_audits
WHERE id LIKE '%provider-secret%'
   OR service_installation_extension_id LIKE '%provider-secret%'
   OR extension_name LIKE '%provider-secret%'
   OR domain LIKE '%provider-secret%'
   OR action LIKE '%provider-secret%'
   OR request_hash LIKE '%provider-secret%'`).Scan(&secretBearingRows))
	require.Zero(t, secretBearingRows)
}

func TestInvokeAmbiguousAppRuntimeConnectionFailsClosed(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_ambiguous")
	executor := &recordingInvocationEgress{}
	service := New(
		db,
		fixture.Registry,
		[]byte(testCredentialKey),
		Options{
			EgressFactory: func(
				registry.ActionPlan,
			) (InvocationEgress, error) {
				return executor, nil
			},
		},
	)
	host, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"invoke-ambiguous",
		fixture.ToolsDigest,
		"tools-primary",
		"agent-1",
		map[string]any{
			"accessToken": "provider-secret-a",
			"scopes":      []string{},
		},
	)
	encoded, err := json.Marshal(map[string]any{
		"accessToken": "provider-secret-b",
		"scopes":      []string{},
	})
	require.NoError(t, err)
	encrypted, err := app.EncryptStoredCredentialsJSON(
		[]byte(testCredentialKey),
		encoded,
	)
	require.NoError(t, err)
	_, err = serviceactor.New(db).UpsertConnection(
		ctx,
		serviceactor.ConnectionUpsert{
			ID:                    "conn_invoke_ambiguous_second",
			WorkspaceID:           fixture.WorkspaceID,
			ServiceDefinitionID:   host.DefinitionID,
			ServiceInstallationID: host.InstallationID,
			ServiceKind:           serviceactor.KindApp,
			ExternalSubjectRef:    "subject-second",
			Credentials:           encrypted,
			Status:                serviceactor.StatusActive,
			CreatedByMemberID:     fixture.HumanMemberID,
			Now:                   time.Now().UTC(),
		},
	)
	require.NoError(t, err)

	_, err = service.Invoke(
		ctx,
		invocationParams(
			fixture,
			"agent-1",
			"tools-primary",
			"list-items",
			fixture.ToolsDigest,
			version,
		),
	)
	requireInvocationCode(t, err, InvocationProviderAuthUnavailable)
	require.Zero(t, executor.callCount())
}

func TestInvokeOAuth401RefreshesAndReplaysExactlyOnce(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_refresh")
	now := time.Date(2026, 8, 11, 10, 0, 0, 0, time.UTC)
	executor := &recordingInvocationEgress{
		results: []egress.Result{
			{
				UpstreamStatus: 401,
				StatusClass:    "4xx",
				ResponseBytes:  2,
			},
			{
				Data:           json.RawMessage(`{"ok":true}`),
				UpstreamStatus: 200,
				StatusClass:    "2xx",
				ResponseBytes:  11,
			},
		},
		errors: []error{
			&egress.UpstreamStatusError{Status: 401},
			nil,
		},
	}
	refresher := &recordingOAuthRefresher{result: OAuthCredential{
		AccessToken:  "fresh-secret",
		RefreshToken: "fresh-refresh",
		Scopes:       []string{},
		ExpiresAt:    now.Add(time.Hour),
	}}
	service := New(
		db,
		fixture.Registry,
		[]byte(testCredentialKey),
		Options{
			Now:            func() time.Time { return now },
			OAuthRefresher: refresher,
			EgressFactory: func(
				registry.ActionPlan,
			) (InvocationEgress, error) {
				return executor, nil
			},
		},
	)
	host, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"invoke-refresh",
		fixture.ToolsDigest,
		"tools-primary",
		"agent-1",
		map[string]any{
			"accessToken":  "stale-secret",
			"refreshToken": "old-refresh",
			"scopes":       []string{},
			"expiresAt":    now.Add(time.Hour).Format(time.RFC3339),
		},
	)

	_, err := service.Invoke(
		ctx,
		invocationParams(
			fixture,
			"agent-1",
			"tools-primary",
			"list-items",
			fixture.ToolsDigest,
			version,
		),
	)
	require.NoError(t, err)
	require.Equal(t, 2, executor.callCount())
	require.Equal(t, "stale-secret", executor.credentials[0].Value)
	require.Equal(t, "fresh-secret", executor.credentials[1].Value)
	require.Equal(t, 1, refresher.calls)

	var stored []byte
	require.NoError(t, db.QueryRow(
		ctx,
		`SELECT credentials FROM service_connections WHERE id=$1`,
		host.ConnectionID,
	).Scan(&stored))
	plaintext, err := app.DecryptStoredCredentialsJSON(
		[]byte(testCredentialKey),
		stored,
	)
	require.NoError(t, err)
	require.Contains(t, string(plaintext), "fresh-secret")
	require.NotContains(t, string(plaintext), "stale-secret")
}

func TestInvokeRejectsMissingScopeAndPlaintextCredentialRows(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_auth_fences")
	pkg := toolsExternalPackage(t)
	pkg.Actions[0].Auth.RequiredScopes = []string{"repo"}
	pkg = finalizePackage(t, pkg)
	reg := registry.New()
	registered, err := reg.RegisterExternal(pkg)
	require.NoError(t, err)
	fixture.Registry = reg
	executor := &recordingInvocationEgress{}
	service := New(
		db,
		reg,
		[]byte(testCredentialKey),
		Options{
			EgressFactory: func(
				registry.ActionPlan,
			) (InvocationEgress, error) {
				return executor, nil
			},
		},
	)
	host, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"auth-fences",
		registered.Definition.Digest,
		"tools-primary",
		"agent-1",
		map[string]any{
			"accessToken": "secret",
			"scopes":      []string{"issues"},
		},
	)
	invoke := invocationParams(
		fixture,
		"agent-1",
		"tools-primary",
		"list-items",
		registered.Definition.Digest,
		version,
	)

	_, err = service.Invoke(ctx, invoke)
	requireInvocationCode(t, err, InvocationProviderAuthUnavailable)
	require.Zero(t, executor.callCount())

	_, err = db.Exec(ctx, `
UPDATE service_connections
SET credentials=$2
WHERE id=$1`,
		host.ConnectionID,
		[]byte(`{"accessToken":"plaintext-secret","scopes":["repo"]}`),
	)
	require.NoError(t, err)
	_, err = service.Invoke(ctx, invoke)
	requireInvocationCode(t, err, InvocationProviderAuthUnavailable)
	require.Zero(t, executor.callCount())
}

func TestInvokeDoesNotReplayMutationOnOrdinaryNetworkFailure(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_no_replay")
	executor := &recordingInvocationEgress{
		results: []egress.Result{{}},
		errors:  []error{egress.ErrRequestFailed},
	}
	service := New(
		db,
		fixture.Registry,
		[]byte(testCredentialKey),
		Options{
			EgressFactory: func(
				registry.ActionPlan,
			) (InvocationEgress, error) {
				return executor, nil
			},
		},
	)
	_, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"no-replay",
		fixture.ToolsDigest,
		"tools-primary",
		"agent-1",
		map[string]any{
			"accessToken": "secret",
			"scopes":      []string{},
		},
	)
	invoke := invocationParams(
		fixture,
		"agent-1",
		"tools-primary",
		"create-item",
		fixture.ToolsDigest,
		version,
	)
	invoke.Input = json.RawMessage(`{"name":"one"}`)

	_, err := service.Invoke(ctx, invoke)
	requireInvocationCode(t, err, InvocationUpstreamFailure)
	require.Equal(
		t,
		1,
		executor.callCount(),
		"ordinary network failures must never replay a mutation",
	)
}

func requireInvocationCode(
	t *testing.T,
	err error,
	expected InvocationErrorCode,
) {
	t.Helper()
	require.Error(t, err)
	code, ok := InvocationErrorCodeOf(err)
	require.True(t, ok)
	require.Equal(t, expected, code)
}

func readAgentCatalogVersion(
	t *testing.T,
	ctx context.Context,
	db coredb.QueryRower,
	workspaceID string,
	agentID string,
) int64 {
	t.Helper()
	var version int64
	require.NoError(t, db.QueryRow(ctx, `
SELECT version FROM agent_extension_catalogs
WHERE workspace_id=$1 AND agent_id=$2`,
		workspaceID,
		agentID,
	).Scan(&version))
	return version
}