package extensionflow

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

	"github.com/stretchr/testify/require"

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

type recordingBuiltinExecutor struct {
	mu               sync.Mutex
	invocations      []registry.BuiltinInvocation
	data             json.RawMessage
	statusClass      string
	auditFailureData json.RawMessage
}

func (executor *recordingBuiltinExecutor) Strategy() string {
	return "recording-native"
}

func (executor *recordingBuiltinExecutor) Invoke(
	_ context.Context,
	invocation registry.BuiltinInvocation,
) (registry.BuiltinResult, error) {
	executor.mu.Lock()
	defer executor.mu.Unlock()
	executor.invocations = append(executor.invocations, invocation)
	data := executor.data
	if len(data) == 0 {
		data = json.RawMessage(`{"builtin":true}`)
	}
	return registry.BuiltinResult{
		Data:                 append(json.RawMessage(nil), data...),
		StatusClass:          executor.statusClass,
		ResultClassification: executor.statusClass,
		AuditFailureData: append(
			json.RawMessage(nil),
			executor.auditFailureData...,
		),
	}, nil
}

type auditFailingDB struct {
	coredb.DB
	auditWrites int
}

func (db *auditFailingDB) Exec(
	ctx context.Context,
	query string,
	args ...any,
) (coredb.Result, error) {
	if strings.Contains(query, "INSERT INTO extension_invocation_audits") {
		db.auditWrites++
		return nil, errors.New("audit unavailable")
	}
	return db.DB.Exec(ctx, query, args...)
}

func TestInvokeBuiltinUsesSameAuthoritativePreflightAndAudit(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_builtin")
	executor := &recordingBuiltinExecutor{statusClass: "unknown"}
	reg := registry.New()
	registered, err := reg.RegisterBuiltin(
		builtinToolsDefinition(),
		executor,
		"recording-v1",
	)
	require.NoError(t, err)
	fixture.Registry = reg
	service := New(
		db,
		reg,
		[]byte(testCredentialKey),
		Options{},
	)
	host, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"invoke-builtin",
		registered.Definition.Digest,
		"builtin-primary",
		"agent-1",
		map[string]any{"appToken": "not-exposed"},
	)

	params := invocationParams(
		fixture,
		"agent-1",
		"builtin-primary",
		"list-items",
		registered.Definition.Digest,
		version,
	)
	preflight, err := service.resolveInvocation(ctx, params)
	require.NoError(t, err)
	require.Regexp(
		t,
		`^sha256:[0-9a-f]{64}$`,
		preflight.connectionGeneration,
	)
	result, err := service.Invoke(ctx, params)
	require.NoError(t, err)
	require.JSONEq(t, `{"builtin":true}`, string(result.Data))
	executor.mu.Lock()
	require.Len(t, executor.invocations, 1)
	invocation := executor.invocations[0]
	executor.mu.Unlock()
	require.Equal(t, fixture.WorkspaceID, invocation.WorkspaceID)
	require.Equal(t, "agent-1", invocation.AgentID)
	require.Equal(
		t,
		"mem_flow_agent_agent-1",
		invocation.AgentMemberID,
	)
	require.Equal(t, host.InstallationID, invocation.ServiceInstallationID)
	require.Equal(t, host.ConnectionID, invocation.ServiceConnectionID)
	require.Equal(
		t,
		preflight.connectionGeneration,
		invocation.ServiceConnectionGeneration,
	)
	require.Equal(
		t,
		host.Selection.ID,
		invocation.ServiceInstallationExtensionID,
	)
	require.Empty(t, invocation.AuthorizedByMemberID)
	require.Equal(t, "builtin-primary", invocation.ExtensionName)
	require.Equal(t, "list-items", invocation.Action)
	require.Equal(
		t,
		registered.Definition.Version,
		invocation.DefinitionVersion,
	)
	require.Equal(
		t,
		registered.Definition.Digest,
		invocation.DefinitionDigest,
	)
	require.Equal(t, version, invocation.CatalogVersion)
	require.Regexp(t, `^sha256:[0-9a-f]{64}$`, invocation.RequestHash)

	var statusClass, classification, selectionID string
	require.NoError(t, db.QueryRow(ctx, `
SELECT execution_status_class, result_classification,
       service_installation_extension_id
FROM extension_invocation_audits
WHERE workspace_id=$1 AND agent_id=$2`,
		fixture.WorkspaceID,
		"agent-1",
	).Scan(&statusClass, &classification, &selectionID))
	require.Equal(t, "unknown", statusClass)
	require.Equal(t, "unknown", classification)
	require.Equal(t, host.Selection.ID, selectionID)
}

func TestInvokeBuiltinMutationAuditFailureReturnsUnknownInsteadOfRetryable500(
	t *testing.T,
) {
	ctx := context.Background()
	realDB := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, realDB, "invoke_audit_unknown")
	executor := &recordingBuiltinExecutor{
		data: json.RawMessage(
			`{"status":"accepted","draft":"opaque-draft"}`,
		),
		statusClass: "success",
		auditFailureData: json.RawMessage(
			`{"status":"unknown"}`,
		),
	}
	reg := registry.New()
	registered, err := reg.RegisterBuiltin(
		builtinToolsDefinition(),
		executor,
		"recording-v1",
	)
	require.NoError(t, err)
	fixture.Registry = reg
	db := &auditFailingDB{DB: realDB}
	service := New(db, reg, []byte(testCredentialKey), Options{})
	_, version := seedAppHostedInvocation(
		t,
		ctx,
		db,
		fixture,
		service,
		"invoke-audit-unknown",
		registered.Definition.Digest,
		"builtin-primary",
		"agent-1",
		map[string]any{"appToken": "not-exposed"},
	)

	result, err := service.Invoke(
		ctx,
		invocationParams(
			fixture,
			"agent-1",
			"builtin-primary",
			"list-items",
			registered.Definition.Digest,
			version,
		),
	)
	require.NoError(t, err)
	require.JSONEq(t, `{"status":"unknown"}`, string(result.Data))
	require.Equal(t, 2, db.auditWrites)
	executor.mu.Lock()
	require.Len(t, executor.invocations, 1)
	executor.mu.Unlock()
	var auditCount int
	require.NoError(t, realDB.QueryRow(ctx, `
SELECT COUNT(*) FROM extension_invocation_audits
WHERE workspace_id=$1 AND agent_id=$2`,
		fixture.WorkspaceID,
		"agent-1",
	).Scan(&auditCount))
	require.Zero(t, auditCount)
}

func TestInvokeConnectionlessBuiltinSkipsProviderCredentials(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "invoke_connectionless_builtin")
	executor := &recordingBuiltinExecutor{}
	reg := registry.New()
	definition := builtinToolsDefinition()
	definition.Domain = "connectionless"
	registered, err := reg.RegisterBuiltinWithOptions(
		definition,
		executor,
		"connectionless-v1",
		registry.BuiltinOptions{
			ConnectionPolicy: registry.ConnectionNone,
		},
	)
	require.NoError(t, err)
	fixture.Registry = reg
	service := New(db, reg, []byte(testCredentialKey), Options{})
	host := seedAppHostedCatalogSelection(
		t,
		ctx,
		db,
		fixture,
		"invoke-connectionless-builtin",
		registered.Definition.Digest,
		"connectionless",
	)
	mustExecService(t, db, `
UPDATE service_connections
SET status='disabled', disabled_at=$2, updated_at=$2
WHERE id=$1`, host.ConnectionID, time.Now().UTC())
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.ApplyAgentCatalogChange(
			ctx,
			tx,
			fixture.WorkspaceID,
			"agent-1",
			time.Now().UTC(),
		)
	}))
	version := readAgentCatalogVersion(
		t,
		ctx,
		db,
		fixture.WorkspaceID,
		"agent-1",
	)

	params := invocationParams(
		fixture,
		"agent-1",
		"connectionless",
		"list-items",
		registered.Definition.Digest,
		version,
	)
	result, err := service.Invoke(ctx, params)
	require.NoError(t, err)
	require.JSONEq(t, `{"builtin":true}`, string(result.Data))
	executor.mu.Lock()
	require.Len(t, executor.invocations, 1)
	invocation := executor.invocations[0]
	executor.mu.Unlock()
	require.Equal(t, host.InstallationID, invocation.ServiceInstallationID)

	resolved, err := service.resolveInvocation(ctx, params)
	require.NoError(t, err)
	require.False(t, resolved.hasConnection)
	require.Empty(t, resolved.connection.ID)
	require.Empty(t, resolved.encryptedCredentials)
}