package extensionflow

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

	"github.com/stretchr/testify/require"

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

var extensionServiceTiDB testkit.TiDBCoreDBFixture

func TestMain(m *testing.M) {
	os.Exit(testkit.RunTestsWithTiDBCoreDBFixture(m, "extension_service", &extensionServiceTiDB))
}

func startServiceRollbackDB(t *testing.T, ctx context.Context) coredb.DB {
	t.Helper()
	return extensionServiceTiDB.SerializedRollbackDB(t, ctx)
}

func startServiceConcurrentDB(t *testing.T, ctx context.Context) coredb.DB {
	t.Helper()
	rootDSN := extensionServiceTiDB.RootDSN(t)
	dsn := testkit.StartInitializedTiDBWithRootDSN(t, rootDSN)
	sqlDB, err := testkit.OpenTiDBContext(ctx, dsn)
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(32)
	sqlDB.SetMaxIdleConns(32)
	t.Cleanup(func() { _ = sqlDB.Close() })
	return coredb.WrapSQLDB(sqlDB)
}

// serviceFixture holds the durable FK rows an Extension control-plane operation
// needs (workspace + human member) plus the resolved Definition digests a test
// registry serves.
type serviceFixture struct {
	WorkspaceID   string
	HumanMemberID string
	Registry      *registry.Registry
	// ToolsDigest identifies an External Definition with read/mutation/
	// destructive actions and a "default" oauth2 auth-profile slot.
	ToolsDigest string
	// SafeDigest identifies an External Definition with only a read action (no
	// destructive) and a distinct domain for extensionName-conflict tests.
	SafeDigest string
	// BuiltinDigest identifies a Builtin Definition used to prove the control
	// plane's kind split (stage-5 round-2): Connect rejects a Builtin (no
	// auth-profile slot), while Grant accepts it (a Builtin carries actions).
	BuiltinDigest string
	// DefaultAuthProfile is the auth-profile slot id declared by the External
	// Tools/Notes Definitions; Connect tests bind credentials to it.
	DefaultAuthProfile string
}

const testCredentialKey = "0123456789abcdef0123456789abcdef" // exactly 32 bytes

func seedServiceFixture(t *testing.T, ctx context.Context, db coredb.QueryExecer, suffix string) serviceFixture {
	t.Helper()
	now := time.Now().UTC()
	orgID := "org_flow_" + suffix
	userID := "usr_flow_" + suffix
	workspaceID := "wsp_flow_" + suffix
	humanMemberID := "mem_flow_human_" + suffix
	mustExecService(t, db, `
INSERT INTO users (id, email, name, display_name, password_hash, created_at, updated_at)
VALUES ($1, $2, $3, $3, 'hash', $4, $4)
`, userID, suffix+"@example.com", suffix, now)
	mustExecService(t, db, `
INSERT INTO organizations (id, slug, name, created_by_user_id, created_at, updated_at)
VALUES ($1, $2, $2, $3, $4, $4)
`, orgID, suffix, userID, now)
	mustExecService(t, db, `
INSERT INTO workspaces (id, organization_id, slug, name, created_at, updated_at)
VALUES ($1, $2, $3, $3, $4, $4)
`, workspaceID, orgID, suffix, now)
	mustExecService(t, db, `
INSERT INTO members (id, workspace_id, type, display_name, user_id, workspace_role, status, created_at, updated_at)
VALUES ($1, $2, 'human', 'Human', $3, 'owner', 'active', $4, $4)
`, humanMemberID, workspaceID, userID, now)

	// Materialize the test agent roster as real workspace agents (member +
	// machine + agent rows) so the Grant agent-identity fence (stage-5 round-3)
	// sees valid in-workspace principals. The roster is the closed set of agent
	// ids the grant/revoke tests authorize.
	for _, agentID := range testAgentRoster {
		seedWorkspaceAgent(t, db, workspaceID, userID, agentID, now)
	}

	reg := registry.New()
	tools, err := reg.RegisterExternal(toolsExternalPackage(t))
	require.NoError(t, err)
	notes, err := reg.RegisterExternal(notesExternalPackage(t))
	require.NoError(t, err)
	builtin, err := reg.RegisterBuiltin(builtinToolsDefinition(), testExecutor{}, "rev-1")
	require.NoError(t, err)

	return serviceFixture{
		WorkspaceID:        workspaceID,
		HumanMemberID:      humanMemberID,
		Registry:           reg,
		ToolsDigest:        tools.Definition.Digest,
		SafeDigest:         notes.Definition.Digest,
		BuiltinDigest:      builtin.Definition.Digest,
		DefaultAuthProfile: testAuthProfile,
	}
}

// testAgentRoster is the closed set of agent ids the service/extension tests
// grant/revoke against; seedServiceFixture materializes each as a real workspace
// agent.
var testAgentRoster = []string{
	"agent-1", "agent-2", "agent-distinct", "agent-never-granted", "agent-ok", "agent-race",
}

// seedWorkspaceAgent materializes one agent principal in the workspace: a
// type='agent' member, a machine, and an agents row. It mirrors the canonical
// agent-resolution shape (members.type='agent' joined to agents) so the Grant
// agent-identity fence (stage-5 round-3) accepts it.
func seedWorkspaceAgent(t *testing.T, db coredb.QueryExecer, workspaceID, userID, agentID string, now time.Time) {
	t.Helper()
	memberID := "mem_flow_agent_" + agentID
	machineID := "mach_flow_agent_" + agentID
	mustExecService(t, db, `
INSERT INTO members (id, workspace_id, type, display_name, status, created_at, updated_at)
VALUES ($1, $2, 'agent', $3, 'active', $4, $4)
`, memberID, workspaceID, agentID, now)
	mustExecService(t, db, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, status, connected_at, disconnected_at, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, 'active', $6, NULL, $6, $6)
`, machineID, workspaceID, userID, agentID, "api_key_"+agentID, now)
	mustExecService(t, db, `
INSERT INTO agents (id, workspace_id, machine_id, member_id, name, runtime, model, description, status, activity, activity_detail, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, 'claude', 'claude', '', 'active', 'active', '', $6, $6)
`, agentID, workspaceID, machineID, memberID, agentID, now)
}

// newService builds a Service over db with the fixture registry and a valid
// 32-byte credential key.
func newService(db coredb.DB, reg *registry.Registry) *Service {
	return New(db, reg, []byte(testCredentialKey), Options{})
}

// testAuthProfile is the single auth-profile slot id declared by the External
// test Definitions; Connect tests bind credentials to it.
const testAuthProfile = "default"

// testEndpoint is the fixed HTTPS origin carried by the External test packages.
const testEndpoint = "https://api.ext-test.example.com"

// objectSchema is the minimal Draft 2020-12 input schema every Definition
// action must declare to pass catalog validation.
func objectSchema() map[string]any {
	return map[string]any{
		"$schema":              "https://json-schema.org/draft/2020-12/schema",
		"type":                 "object",
		"additionalProperties": true,
	}
}

// passthroughResponse is the deterministic passthrough response binding shared
// by every External test action.
func passthroughResponse() spec.Response {
	return spec.Response{
		Mode:            "passthrough",
		SuccessStatuses: []int{200},
		ContentTypes:    []string{"application/json"},
	}
}

// authProfiles returns the single "default" oauth2 auth-profile slot shared by
// the External test packages.
func authProfiles() spec.Auth {
	return spec.Auth{
		Default: testAuthProfile,
		Profiles: []spec.AuthProfile{
			{ID: testAuthProfile, Type: "oauth2_bearer"},
		},
	}
}

func requireAction(name, summary string, safety spec.Safety, method, path string) spec.Action {
	return spec.Action{
		Name: name,
		File: "actions/" + name + ".yaml",
		Auth: &spec.ActionAuthRequirement{Profile: testAuthProfile},
		Catalog: spec.Catalog{
			Summary:      summary,
			Description:  summary,
			Safety:       safety,
			InputSchema:  objectSchema(),
			OutputSchema: objectSchema(),
		},
		Binding: spec.Binding{
			Method:   method,
			Path:     path,
			Response: passthroughResponse(),
		},
	}
}

// toolsExternalPackage builds a self-consistent External Definition package with
// one read, one mutation, and one destructive action. The destructive action is
// what the explicit-consent invariant exercises.
func toolsExternalPackage(t *testing.T) spec.Package {
	t.Helper()
	pkg := spec.Package{
		Metadata: spec.Metadata{
			SchemaVersion: 1,
			Domain:        "tools",
			Version:       1,
			Name:          "Tools",
			Description:   "Tools actions curated for tests.",
			Endpoint:      spec.Endpoint{BaseURL: testEndpoint},
			Auth:          authProfiles(),
		},
		Actions: []spec.Action{
			requireAction("list-items", "List", spec.Safety{Class: "read", Idempotent: true}, "GET", "/items"),
			requireAction("create-item", "Create", spec.Safety{Class: "mutation"}, "POST", "/items"),
			requireAction("purge-item", "Purge", spec.Safety{Class: "destructive", Destructive: true}, "DELETE", "/items"),
		},
	}
	return finalizePackage(t, pkg)
}

// notesExternalPackage builds a self-consistent External Definition whose only
// action is non-destructive; it provides a second domain for extensionName-conflict
// tests.
func notesExternalPackage(t *testing.T) spec.Package {
	t.Helper()
	pkg := spec.Package{
		Metadata: spec.Metadata{
			SchemaVersion: 1,
			Domain:        "notes",
			Version:       1,
			Name:          "Notes",
			Description:   "Notes actions curated for tests.",
			Endpoint:      spec.Endpoint{BaseURL: testEndpoint},
			Auth:          authProfiles(),
		},
		Actions: []spec.Action{
			requireAction("read-note", "Read", spec.Safety{Class: "read", Idempotent: true}, "GET", "/notes"),
		},
	}
	return finalizePackage(t, pkg)
}

// finalizePackage computes and sets the canonical digest so the package is
// self-consistent for RegisterExternal.
func finalizePackage(t *testing.T, pkg spec.Package) spec.Package {
	t.Helper()
	digest, err := spec.DefinitionDigest(pkg)
	require.NoError(t, err)
	pkg.Metadata.Digest = digest
	pkg.Digest = digest
	return pkg
}

// builtinToolsDefinition is a COMMON Definition (no HTTP binding, no auth
// profiles) registered as a Builtin. It is used to prove the control plane's
// kind split (stage-5 round-2): Grant accepts a Builtin installation (it carries
// actions), while Connect rejects it (a Builtin has no auth-profile slot).
func builtinToolsDefinition() registry.Definition {
	return registry.Definition{
		Domain:      "builtin-tools",
		Version:     1,
		Name:        "Builtin Tools",
		Description: "Builtin tools actions curated for wrong-kind tests.",
		Actions: []registry.DefinitionAction{
			{Name: "list-items", Summary: "List", Safety: spec.Safety{Class: "read", Idempotent: true}, InputSchema: objectSchema()},
		},
	}
}

func registerProviderWatchBuiltinsForTest(reg *registry.Registry) error {
	for _, domain := range []string{"github", "gitea"} {
		definition := registry.Definition{
			Domain:      domain,
			Version:     1,
			Name:        domain + " pull request watches",
			Description: "Test App-hosted provider watch surface.",
			Actions: []registry.DefinitionAction{
				{Name: "watch-pull-request", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
				{Name: "list-watches", Safety: spec.Safety{Class: "read", Idempotent: true}, InputSchema: objectSchema()},
				{Name: "pause-watch", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
				{Name: "resume-watch", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
				{Name: "unwatch", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
			},
		}
		if _, err := reg.RegisterBuiltin(
			definition,
			testExecutor{},
			"provider-watch-test-v1",
		); err != nil {
			return err
		}
	}
	return nil
}

// testExecutor is a no-op Builtin executor used only to satisfy the registry's
// Builtin registration contract in tests.
type testExecutor struct{}

func (testExecutor) Strategy() string { return "test" }

func (testExecutor) Invoke(context.Context, registry.BuiltinInvocation) (registry.BuiltinResult, error) {
	return registry.BuiltinResult{Data: json.RawMessage(`{"ok":true}`)}, nil
}

func mustExecService(t testing.TB, execer coredb.QueryExecer, query string, args ...any) {
	t.Helper()
	_, err := execer.Exec(context.Background(), query, args...)
	require.NoError(t, err)
}

// registerProviderWatchBuiltinsWithPolicyForTest registers the GitHub test
// Builtin with a provider catalog policy facet attached ATOMICALLY through
// BuiltinOptions — the only supported wiring since runtime policy
// attachment was removed. The digest matches the plain variant (same
// definition, revision, and ConnectionRequired policy).
func registerProviderWatchBuiltinsWithPolicyForTest(
	reg *registry.Registry,
	policy registry.CatalogPolicy,
) error {
	for _, domain := range []string{"github", "gitea"} {
		definition := registry.Definition{
			Domain:      domain,
			Version:     1,
			Name:        domain + " pull request watches",
			Description: "Test App-hosted provider watch surface.",
			Actions: []registry.DefinitionAction{
				{Name: "watch-pull-request", Summary: "Watch", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
				{Name: "list-watches", Summary: "List", Safety: spec.Safety{Class: "read", Idempotent: true}, InputSchema: objectSchema()},
				{Name: "pause-watch", Summary: "Pause", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
				{Name: "resume-watch", Summary: "Resume", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
				{Name: "unwatch", Summary: "Unwatch", Safety: spec.Safety{Class: "mutation"}, InputSchema: objectSchema()},
			},
		}
		var facet registry.CatalogPolicy
		if domain == "github" {
			facet = policy
		}
		if _, err := reg.RegisterBuiltinWithOptions(
			definition,
			testExecutor{},
			"provider-watch-test-v1",
			registry.BuiltinOptions{
				ConnectionPolicy: registry.ConnectionRequired,
				CatalogPolicy:    facet,
			},
		); err != nil {
			return err
		}
	}
	return nil
}