package extensionflow

import (
	"context"
	"encoding/json"
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func seedAppHostedInvocation(
	t *testing.T,
	ctx context.Context,
	db coredb.DB,
	fixture serviceFixture,
	service *Service,
	suffix string,
	definitionDigest string,
	extensionName string,
	agentID string,
	credentials map[string]any,
) (appHostedCatalogFixture, int64) {
	t.Helper()
	host := seedAppHostedCatalogSelection(
		t,
		ctx,
		db,
		fixture,
		suffix,
		definitionDigest,
		extensionName,
	)
	encoded, err := json.Marshal(credentials)
	require.NoError(t, err)
	encrypted, err := app.EncryptStoredCredentialsJSON(
		[]byte(testCredentialKey),
		encoded,
	)
	require.NoError(t, err)
	_, err = serviceactor.New(db).UpsertConnection(
		ctx,
		serviceactor.ConnectionUpsert{
			ID:                    host.ConnectionID,
			WorkspaceID:           fixture.WorkspaceID,
			ServiceDefinitionID:   host.DefinitionID,
			ServiceInstallationID: host.InstallationID,
			ServiceKind:           serviceactor.KindApp,
			ExternalSubjectRef:    "subject-" + suffix,
			Credentials:           encrypted,
			Status:                serviceactor.StatusActive,
			CreatedByMemberID:     fixture.HumanMemberID,
			Now:                   time.Now().UTC(),
		},
	)
	require.NoError(t, err)
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.ApplyAgentCatalogChange(
			ctx,
			tx,
			fixture.WorkspaceID,
			agentID,
			time.Now().UTC(),
		)
	}))
	return host, readAgentCatalogVersion(
		t,
		ctx,
		db,
		fixture.WorkspaceID,
		agentID,
	)
}

func invocationParams(
	fixture serviceFixture,
	agentID string,
	extensionName string,
	action string,
	digest string,
	version int64,
) InvokeParams {
	return InvokeParams{
		WorkspaceID:   fixture.WorkspaceID,
		AgentID:       agentID,
		AgentMemberID: "mem_flow_agent_" + agentID,
		ExtensionName: extensionName,
		Action:        action,
		Input:         json.RawMessage(`{}`),
	}
}

func TestAppHostedInvocationResolutionSuccess(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_resolution_success")
	service := newService(db, fixture.Registry)
	host, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"invoke-resolution-success",
		fixture.ToolsDigest,
		"workspace-tools",
		"agent-1",
		map[string]any{
			"accessToken": "provider-secret",
			"scopes":      []string{},
		},
	)

	resolved, err := service.resolveInvocation(
		ctx,
		invocationParams(
			fixture,
			"agent-1",
			"workspace-tools",
			"list-items",
			fixture.ToolsDigest,
			version,
		),
	)
	require.NoError(t, err)
	require.Equal(t, host.Selection.ID, resolved.selection.ID)
	require.Equal(t, host.InstallationID, resolved.selection.ServiceInstallationID)
	require.Equal(t, fixture.ToolsDigest, resolved.definition.Definition.Digest)
	require.Equal(t, "list-items", resolved.action.Name)
	require.Equal(t, registry.External, resolved.kind)
	require.Equal(t, serviceactor.KindApp, resolved.connection.ServiceKind)
	require.NotContains(t, string(resolved.encryptedCredentials), "provider-secret")
}

func TestAppHostedInvocationResolutionRejectsBeforeExecution(t *testing.T) {
	tests := []struct {
		name     string
		mutate   func(t *testing.T, db coredb.DB, fixture serviceFixture, host appHostedCatalogFixture)
		params   func(base InvokeParams) InvokeParams
		wantCode InvocationErrorCode
	}{
		{
			name: "authenticated member mismatch",
			params: func(base InvokeParams) InvokeParams {
				base.AgentMemberID = "foreign-member"
				return base
			},
			wantCode: InvocationPermissionDenied,
		},
		{
			name: "unknown flattened name",
			params: func(base InvokeParams) InvokeParams {
				base.ExtensionName = "missing-tools"
				return base
			},
			wantCode: InvocationNotFound,
		},
		{
			name: "unknown action",
			params: func(base InvokeParams) InvokeParams {
				base.Action = "missing-action"
				return base
			},
			wantCode: InvocationNotFound,
		},
		{
			name: "disabled selection",
			mutate: func(
				t *testing.T,
				db coredb.DB,
				fixture serviceFixture,
				host appHostedCatalogFixture,
			) {
				_, err := serviceactor.New(db).
					DisableInstallationExtension(
						context.Background(),
						fixture.WorkspaceID,
						host.Selection.ID,
						"test",
						time.Now().UTC(),
					)
				require.NoError(t, err)
			},
			wantCode: InvocationAgentCatalogStale,
		},
		{
			name: "disabled App installation",
			mutate: func(
				t *testing.T,
				db coredb.DB,
				_ serviceFixture,
				host appHostedCatalogFixture,
			) {
				_, err := db.Exec(context.Background(), `
UPDATE service_installations
SET status='disabled', disabled_at=$2, updated_at=$2
WHERE id=$1`, host.InstallationID, time.Now().UTC())
				require.NoError(t, err)
			},
			wantCode: InvocationInternal,
		},
		{
			name: "disabled App definition",
			mutate: func(
				t *testing.T,
				db coredb.DB,
				_ serviceFixture,
				host appHostedCatalogFixture,
			) {
				_, err := db.Exec(context.Background(), `
UPDATE service_definitions
SET status='disabled', updated_at=$2
WHERE id=$1`, host.DefinitionID, time.Now().UTC())
				require.NoError(t, err)
			},
			wantCode: InvocationInternal,
		},
		{
			name: "disabled App Extension capability",
			mutate: func(
				t *testing.T,
				db coredb.DB,
				_ serviceFixture,
				host appHostedCatalogFixture,
			) {
				_, err := db.Exec(context.Background(), `
UPDATE service_capabilities
SET default_enabled=FALSE, updated_at=$2
WHERE service_definition_id=$1
  AND capability_key='extensions'`,
					host.DefinitionID,
					time.Now().UTC(),
				)
				require.NoError(t, err)
			},
			wantCode: InvocationInternal,
		},
		{
			name: "disabled App connection",
			mutate: func(
				t *testing.T,
				db coredb.DB,
				_ serviceFixture,
				host appHostedCatalogFixture,
			) {
				_, err := db.Exec(context.Background(), `
UPDATE service_connections
SET status='disabled', disabled_at=$2, updated_at=$2
WHERE id=$1`, host.ConnectionID, time.Now().UTC())
				require.NoError(t, err)
			},
			wantCode: InvocationAgentCatalogStale,
		},
		{
			name: "disabled App authorization",
			mutate: func(
				t *testing.T,
				db coredb.DB,
				fixture serviceFixture,
				host appHostedCatalogFixture,
			) {
				authorizeAppHostForAgent(
					t,
					context.Background(),
					db,
					fixture,
					host,
					"agent-1",
					false,
				)
			},
			wantCode: InvocationAgentCatalogStale,
		},
	}
	for index, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			ctx := context.Background()
			db := startServiceRollbackDB(t, ctx)
			suffix := "invoke_resolution_" + strconv.Itoa(index)
			fixture := seedServiceFixture(t, ctx, db, suffix)
			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,
				suffix,
				fixture.ToolsDigest,
				"workspace-tools",
				"agent-1",
				map[string]any{
					"accessToken": "provider-secret",
					"scopes":      []string{},
				},
			)
			if test.mutate != nil {
				test.mutate(t, db, fixture, host)
			}
			params := invocationParams(
				fixture,
				"agent-1",
				"workspace-tools",
				"list-items",
				fixture.ToolsDigest,
				version,
			)
			if test.params != nil {
				params = test.params(params)
			}
			_, err := service.Invoke(ctx, params)
			requireInvocationCode(t, err, test.wantCode)
			require.Zero(t, executor.callCount())
		})
	}
}