package extension

import (
	"context"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"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/egress"
	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/extension/spec"
	"github.com/tidbcloud/AgentTeam/internal/service/extensionflow"
	"github.com/tidbcloud/AgentTeam/internal/store/app"
	"github.com/tidbcloud/AgentTeam/internal/store/serviceactor"
)

var extensionAPITiDB testkit.TiDBCoreDBFixture

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

func TestInvokeProductionRouteToGuardedTLSAndAudit(t *testing.T) {
	ctx := context.Background()
	db := extensionAPITiDB.RollbackDB(t, ctx)
	now := time.Date(2026, 8, 11, 12, 0, 0, 0, time.UTC)
	seedInvocationAPIFixture(t, db, now)

	var receivedAuth string
	provider := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		receivedAuth = request.Header.Get("Authorization")
		require.Equal(t, http.MethodGet, request.Method)
		require.Equal(t, "/v1/items", request.URL.Path)
		writer.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(writer, `{"items":["one"]}`)
	}))
	defer provider.Close()
	providerOrigin := publicProviderOrigin(t, provider, "api-provider.example")

	reg := registry.New()
	pkg := invocationE2EPackage(t, providerOrigin)
	registered, err := reg.RegisterExternal(pkg)
	require.NoError(t, err)
	key := []byte("0123456789abcdef0123456789abcdef")
	service := extensionflow.New(db, reg, key, extensionflow.Options{
		Now: func() time.Time { return now },
		EgressFactory: func(plan registry.ActionPlan) (extensionflow.InvocationEgress, error) {
			origin, err := url.Parse(plan.Endpoint.BaseURL)
			if err != nil {
				return nil, err
			}
			tlsConfig := provider.Client().Transport.(*http.Transport).
				TLSClientConfig.Clone()
			tlsConfig.ServerName = "example.com"
			return egress.NewClient(egress.Config{
				Origin:          origin,
				Resolver:        publicTestResolver{},
				DialContext:     localServerDialer(provider),
				TLSClientConfig: tlsConfig,
			})
		},
	})
	seedInvocationAPIAppHost(
		t,
		ctx,
		db,
		service,
		registered,
		key,
		now,
	)
	router := newRouterWithInvoker(t, authenticator(), reg.List(), service)
	rootTarget := "/api/v1/workspaces/" + testWorkspace + "/agents/" + testAgent + "/runtime/ext"
	recorder := doRaw(
		t,
		router,
		http.MethodPost,
		rootTarget+"/e2e-provider/actions/list-items",
		testToken,
		mustJSON(t, invokeRequest{
			Input: json.RawMessage(`{}`),
		}),
	)
	require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
	require.Contains(t, recorder.Body.String(), `"items":["one"]`)
	require.Equal(t, "Bearer provider-secret", receivedAuth)

	var audits int
	var classification, requestHash string
	require.NoError(t, db.QueryRow(ctx, `
SELECT COUNT(*), MAX(result_classification), MAX(request_hash)
FROM extension_invocation_audits
WHERE workspace_id=$1 AND agent_id=$2`,
		testWorkspace,
		testAgent,
	).Scan(&audits, &classification, &requestHash))
	require.Equal(t, 1, audits)
	require.Equal(t, "success", classification)
	require.Regexp(t, `^sha256:[0-9a-f]{64}$`, requestHash)
	require.NotContains(t, recorder.Body.String(), "provider-secret")
	require.NotContains(t, recorder.Body.String(), provider.URL)
}

type publicTestResolver struct{}

func (publicTestResolver) LookupIP(
	context.Context,
	string,
	string,
) ([]net.IP, error) {
	return []net.IP{net.ParseIP("8.8.8.8")}, nil
}

func publicProviderOrigin(
	t *testing.T,
	server *httptest.Server,
	host string,
) string {
	t.Helper()
	_, port, err := net.SplitHostPort(server.Listener.Addr().String())
	require.NoError(t, err)
	return (&url.URL{
		Scheme: "https",
		Host:   net.JoinHostPort(host, port),
	}).String()
}

func localServerDialer(
	server *httptest.Server,
) func(context.Context, string, string) (net.Conn, error) {
	address := server.Listener.Addr().String()
	return func(ctx context.Context, network, _ string) (net.Conn, error) {
		return (&net.Dialer{}).DialContext(ctx, network, address)
	}
}

func seedInvocationAPIAppHost(
	t *testing.T,
	ctx context.Context,
	db coredb.DB,
	service *extensionflow.Service,
	registered registry.RegisteredExtension,
	key []byte,
	now time.Time,
) {
	t.Helper()
	const (
		definitionID   = "app_ext_api"
		installationID = "appinst_ext_api"
		connectionID   = "appconn_ext_api"
		selectionID    = "svcinstext_ext_api"
	)
	actors := serviceactor.New(db)
	_, err := actors.UpsertDefinition(ctx, serviceactor.DefinitionUpsert{
		ID:               definitionID,
		ServiceKind:      serviceactor.KindApp,
		InstallationMode: serviceactor.InstallationModeNamedMultiInstance,
		Slug:             "ext-api-provider",
		Name:             "Extension API provider",
		LatestVersion:    "1.0.0",
		Status:           serviceactor.StatusActive,
		Now:              now,
	})
	require.NoError(t, err)
	_, err = actors.UpsertInstallation(ctx, serviceactor.InstallationUpsert{
		ID:                  installationID,
		WorkspaceID:         testWorkspace,
		ServiceDefinitionID: definitionID,
		ServiceKind:         serviceactor.KindApp,
		InstanceKey:         "primary",
		Status:              serviceactor.StatusActive,
		InstalledVersion:    "1.0.0",
		CreatedByMemberID:   "mem_owner",
		Now:                 now,
	})
	require.NoError(t, err)
	capability, err := json.Marshal(map[string]any{
		"extensions": []map[string]any{{
			"domain":           registered.Definition.Domain,
			"version":          registered.Definition.Version,
			"definitionDigest": registered.Definition.Digest,
			"defaultName":      "e2e-provider",
			"defaultEnabled":   true,
		}},
	})
	require.NoError(t, err)
	_, err = actors.UpsertCapability(ctx, serviceactor.CapabilityUpsert{
		ID:                  "svccap_ext_api",
		ServiceDefinitionID: definitionID,
		ServiceKind:         serviceactor.KindApp,
		CapabilityKey:       serviceactor.CapabilityExtensions,
		Config:              capability,
		DefaultEnabled:      true,
		Now:                 now,
	})
	require.NoError(t, err)
	credentials, err := json.Marshal(map[string]any{
		"accessToken": "provider-secret",
		"scopes":      []string{},
	})
	require.NoError(t, err)
	encrypted, err := app.EncryptStoredCredentialsJSON(key, credentials)
	require.NoError(t, err)
	_, err = actors.UpsertConnection(ctx, serviceactor.ConnectionUpsert{
		ID:                    connectionID,
		WorkspaceID:           testWorkspace,
		ServiceDefinitionID:   definitionID,
		ServiceInstallationID: installationID,
		ServiceKind:           serviceactor.KindApp,
		ExternalSubjectRef:    "provider-user",
		Credentials:           encrypted,
		Status:                serviceactor.StatusActive,
		CreatedByMemberID:     "mem_owner",
		Now:                   now,
	})
	require.NoError(t, err)
	_, err = actors.UpsertInstallationExtension(
		ctx,
		serviceactor.InstallationExtensionUpsert{
			InstallationExtension: serviceactor.InstallationExtension{
				ID:                    selectionID,
				WorkspaceID:           testWorkspace,
				ServiceDefinitionID:   definitionID,
				ServiceInstallationID: installationID,
				ExtensionName:         "e2e-provider",
				DefinitionDomain:      registered.Definition.Domain,
				DefinitionVersion:     registered.Definition.Version,
				DefinitionDigest:      registered.Definition.Digest,
				Status:                serviceactor.StatusActive,
				SelectionSource: serviceactor.
					InstallationExtensionSourceUser,
				CreatedByMemberID: "mem_owner",
			},
			Now: now,
		},
	)
	require.NoError(t, err)
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.ApplyAgentCatalogChange(
			ctx,
			tx,
			testWorkspace,
			testAgent,
			now,
		)
	}))
}

func seedInvocationAPIFixture(t *testing.T, db coredb.QueryExecer, now time.Time) {
	t.Helper()
	mustExecInvocationAPI(t, db, `
INSERT INTO users (id, email, name, display_name, password_hash, created_at, updated_at)
VALUES ('usr_ext_api', 'ext-api@example.com', 'Extension API', 'Extension API', 'hash', $1, $1)`,
		now)
	mustExecInvocationAPI(t, db, `
INSERT INTO organizations (id, slug, name, created_by_user_id, created_at, updated_at)
VALUES ('org_ext_api', 'ext-api', 'Extension API', 'usr_ext_api', $1, $1)`,
		now)
	mustExecInvocationAPI(t, db, `
INSERT INTO workspaces (id, organization_id, slug, name, created_at, updated_at)
VALUES ($1, 'org_ext_api', 'ext-api', 'Extension API', $2, $2)`,
		testWorkspace, now)
	mustExecInvocationAPI(t, db, `
INSERT INTO members (
    id, workspace_id, type, display_name, user_id, workspace_role,
    status, created_at, updated_at
) VALUES ('mem_owner', $1, 'human', 'Owner', 'usr_ext_api', 'owner', 'active', $2, $2)`,
		testWorkspace, now)
	mustExecInvocationAPI(t, db, `
INSERT INTO members (id, workspace_id, type, display_name, status, created_at, updated_at)
VALUES ($1, $2, 'agent', 'Agent', 'active', $3, $3)`,
		testMember, testWorkspace, now)
	mustExecInvocationAPI(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, 'usr_ext_api', 'Agent machine', 'api-key', 'active', $3, NULL, $3, $3)`,
		testMachine, testWorkspace, now)
	mustExecInvocationAPI(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, 'Agent', 'claude', 'claude', '', 'active', 'active', '', $5, $5)`,
		testAgent, testWorkspace, testMachine, testMember, now)
}

func invocationE2EPackage(t *testing.T, endpoint string) spec.Package {
	t.Helper()
	schema := map[string]any{
		"$schema":              "https://json-schema.org/draft/2020-12/schema",
		"type":                 "object",
		"additionalProperties": false,
	}
	pkg := spec.Package{
		Metadata: spec.Metadata{
			SchemaVersion: 1, Domain: "e2e-provider", Version: 1,
			Name: "E2E provider", Endpoint: spec.Endpoint{BaseURL: endpoint + "/v1"},
			Auth: spec.Auth{
				Default: "default",
				Profiles: []spec.AuthProfile{{
					ID: "default", Type: "oauth2_bearer",
				}},
			},
		},
		Actions: []spec.Action{{
			Name: "list-items", File: "actions/list-items.yaml",
			Auth: &spec.ActionAuthRequirement{Profile: "default"},
			Catalog: spec.Catalog{
				Summary: "List items", Safety: spec.Safety{
					Class: "read", Idempotent: true,
				},
				InputSchema: schema, OutputSchema: schema,
			},
			Binding: spec.Binding{
				Method: http.MethodGet, Path: "/items",
				Response: spec.Response{
					Mode: "passthrough", SuccessStatuses: []int{http.StatusOK},
					ContentTypes: []string{"application/json"},
				},
			},
		}},
	}
	digest, err := spec.DefinitionDigest(pkg)
	require.NoError(t, err)
	pkg.Metadata.Digest = digest
	pkg.Digest = digest
	return pkg
}

func mustExecInvocationAPI(
	t *testing.T,
	db coredb.QueryExecer,
	query string,
	args ...any,
) {
	t.Helper()
	_, err := db.Exec(context.Background(), query, args...)
	require.NoError(t, err)
}

func mustJSON(t *testing.T, value any) string {
	t.Helper()
	encoded, err := json.Marshal(value)
	require.NoError(t, err)
	return string(encoded)
}