package extensione2e

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/exec"
	"github.com/tidbcloud/AgentTeam/desktop/internal/promptcompiler"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/service/extensionflow"
	"github.com/tidbcloud/AgentTeam/internal/store/app"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
	"github.com/tidbcloud/AgentTeam/internal/store/serviceactor"
)

type appHostedFixture struct {
	selection serviceactor.InstallationExtension
}

func requireAppHostedIdentityOnly(
	t *testing.T,
	ctx context.Context,
	db coredb.DB,
) {
	t.Helper()
	var definitionKind, installationKind, connectionKind, authorizationKind string
	require.NoError(t, db.QueryRow(ctx, `
SELECT definition.service_kind,
       installation.service_kind,
       connection.service_kind,
       authorization.service_kind
FROM service_definitions definition
JOIN service_installations installation
  ON installation.service_definition_id = definition.id
JOIN service_connections connection
  ON connection.service_installation_id = installation.id
JOIN service_principal_authorizations authorization
  ON authorization.service_installation_id = installation.id
WHERE definition.id = $1
  AND installation.id = $2
  AND connection.id = $3
  AND authorization.id = $4`,
		e2eHostDefinitionID,
		e2eHostInstallationID,
		e2eHostConnectionID,
		e2eHostAuthorization,
	).Scan(
		&definitionKind,
		&installationKind,
		&connectionKind,
		&authorizationKind,
	))
	require.Equal(t, serviceactor.KindApp, definitionKind)
	require.Equal(t, serviceactor.KindApp, installationKind)
	require.Equal(t, serviceactor.KindApp, connectionKind)
	require.Equal(t, serviceactor.KindApp, authorizationKind)
	var independentExtensionActors int
	require.NoError(t, db.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM service_definitions
   WHERE service_kind = 'extension') +
  (SELECT COUNT(*) FROM service_installations
   WHERE service_kind = 'extension') +
  (SELECT COUNT(*) FROM service_connections
   WHERE service_kind = 'extension') +
  (SELECT COUNT(*) FROM service_principal_authorizations
   WHERE service_kind = 'extension')`,
	).Scan(&independentExtensionActors))
	require.Zero(t, independentExtensionActors)
}

func seedAppHostedExtension(
	t *testing.T,
	ctx context.Context,
	db coredb.DB,
	service *extensionflow.Service,
	registered registry.RegisteredExtension,
	credentialKey []byte,
	now time.Time,
) appHostedFixture {
	t.Helper()
	actors := serviceactor.New(db)
	_, err := actors.UpsertDefinition(ctx, serviceactor.DefinitionUpsert{
		ID:               e2eHostDefinitionID,
		ServiceKind:      serviceactor.KindApp,
		InstallationMode: serviceactor.InstallationModeNamedMultiInstance,
		Slug:             "extension-full-e2e-app",
		Name:             "Extension Full E2E App",
		LatestVersion:    "1.0.0",
		Status:           serviceactor.StatusActive,
		Now:              now,
	})
	require.NoError(t, err)
	_, err = actors.UpsertInstallation(ctx, serviceactor.InstallationUpsert{
		ID:                  e2eHostInstallationID,
		WorkspaceID:         e2eWorkspace,
		ServiceDefinitionID: e2eHostDefinitionID,
		ServiceKind:         serviceactor.KindApp,
		InstanceKey:         "primary",
		Status:              serviceactor.StatusActive,
		InstalledVersion:    "1.0.0",
		CreatedByMemberID:   "mem_owner_extension_full_e2e",
		Now:                 now,
	})
	require.NoError(t, err)
	capabilityJSON, err := model.EncodeExtensionCapabilityConfig(
		model.ExtensionCapabilityConfig{
			Extensions: []model.ExtensionCapabilityEntry{{
				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:                  "cap_extension_full_e2e",
		ServiceDefinitionID: e2eHostDefinitionID,
		ServiceKind:         serviceactor.KindApp,
		CapabilityKey:       serviceactor.CapabilityExtensions,
		Config:              capabilityJSON,
		DefaultEnabled:      true,
		Now:                 now,
	})
	require.NoError(t, err)
	credentialJSON, err := json.Marshal(map[string]any{
		"accessToken": "provider-secret",
		"scopes":      []string{},
	})
	require.NoError(t, err)
	encryptedCredentials, err := app.EncryptStoredCredentialsJSON(
		credentialKey,
		credentialJSON,
	)
	require.NoError(t, err)
	_, err = actors.UpsertConnection(ctx, serviceactor.ConnectionUpsert{
		ID:                    e2eHostConnectionID,
		WorkspaceID:           e2eWorkspace,
		ServiceDefinitionID:   e2eHostDefinitionID,
		ServiceInstallationID: e2eHostInstallationID,
		ServiceKind:           serviceactor.KindApp,
		ExternalSubjectRef:    e2eHostSubject,
		Credentials:           encryptedCredentials,
		Status:                serviceactor.StatusActive,
		CreatedByMemberID:     "mem_owner_extension_full_e2e",
		Now:                   now,
	})
	require.NoError(t, err)
	selection, err := actors.UpsertInstallationExtension(
		ctx,
		serviceactor.InstallationExtensionUpsert{
			InstallationExtension: serviceactor.InstallationExtension{
				ID:                    e2eHostSelectionID,
				WorkspaceID:           e2eWorkspace,
				ServiceDefinitionID:   e2eHostDefinitionID,
				ServiceInstallationID: e2eHostInstallationID,
				ExtensionName:         "e2e-provider",
				DefinitionDomain:      registered.Definition.Domain,
				DefinitionVersion:     registered.Definition.Version,
				DefinitionDigest:      registered.Definition.Digest,
				Status:                serviceactor.StatusActive,
				SelectionSource:       serviceactor.InstallationExtensionSourceUser,
				CreatedByMemberID:     "mem_owner_extension_full_e2e",
			},
			Now: now,
		},
	)
	require.NoError(t, err)
	_, err = actors.UpsertPrincipalAuthorization(
		ctx,
		serviceactor.PrincipalAuthorizationUpsert{
			ID:                     e2eHostAuthorization,
			WorkspaceID:            e2eWorkspace,
			ServiceDefinitionID:    e2eHostDefinitionID,
			ServiceInstallationID:  e2eHostInstallationID,
			ServiceKind:            serviceactor.KindApp,
			PrincipalKind:          "agent",
			PrincipalRef:           e2eAgent,
			AuthorizedByMemberID:   "mem_owner_extension_full_e2e",
			Now:                    now,
			ActiveAuthorizationKey: e2eAuthorizationKey,
		},
	)
	require.NoError(t, err)
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.ApplyAgentCatalogChange(
			ctx,
			tx,
			e2eWorkspace,
			e2eAgent,
			now,
		)
	}))
	return appHostedFixture{
		selection: selection,
	}
}

func seedSecondActiveAppConnection(
	t *testing.T,
	ctx context.Context,
	db coredb.DB,
	credentialKey []byte,
	now time.Time,
) string {
	t.Helper()
	credentialJSON, err := json.Marshal(map[string]any{
		"accessToken": "provider-secret-second",
		"scopes":      []string{},
	})
	require.NoError(t, err)
	encryptedCredentials, err := app.EncryptStoredCredentialsJSON(
		credentialKey,
		credentialJSON,
	)
	require.NoError(t, err)
	_, err = serviceactor.New(db).UpsertConnection(
		ctx,
		serviceactor.ConnectionUpsert{
			ID:                    e2eHostSecondConnectionID,
			WorkspaceID:           e2eWorkspace,
			ServiceDefinitionID:   e2eHostDefinitionID,
			ServiceInstallationID: e2eHostInstallationID,
			ServiceKind:           serviceactor.KindApp,
			ExternalSubjectRef:    e2eHostSecondSubject,
			Credentials:           encryptedCredentials,
			Status:                serviceactor.StatusActive,
			CreatedByMemberID:     "mem_owner_extension_full_e2e",
			Now:                   now,
		},
	)
	require.NoError(t, err)
	return e2eHostSecondConnectionID
}

const e2eAuthorizationKey = "extension-e2e/agent/app-host"

func disableAppAuthorization(
	t *testing.T,
	ctx context.Context,
	db coredb.DB,
	service *extensionflow.Service,
	now time.Time,
) {
	t.Helper()
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := serviceactor.New(db).InTx(tx).DisablePrincipalAuthorization(
			ctx,
			serviceactor.PrincipalAuthorizationDisable{
				ID:                     e2eHostAuthorization,
				WorkspaceID:            e2eWorkspace,
				ServiceDefinitionID:    e2eHostDefinitionID,
				ServiceInstallationID:  e2eHostInstallationID,
				ServiceKind:            serviceactor.KindApp,
				PrincipalKind:          "agent",
				PrincipalRef:           e2eAgent,
				DisabledReason:         "offline_revoke_e2e",
				AuthorizedByMemberID:   "mem_owner_extension_full_e2e",
				Now:                    now,
				ActiveAuthorizationKey: e2eAuthorizationKey,
			},
		)
		if err != nil {
			return err
		}
		return service.ApplyAgentCatalogChange(
			ctx,
			tx,
			e2eWorkspace,
			e2eAgent,
			now,
		)
	}))
}

func requireFlattenedDesktopProjection(
	t *testing.T,
	snapshot exec.ExtensionSnapshot,
) {
	t.Helper()
	serialized, err := json.Marshal(snapshot.Domains())
	require.NoError(t, err)
	help := snapshot.RenderRootHelp()
	guidance := snapshot.RuntimeGuidance()
	facts := snapshot.PromptGuidanceActions()
	actions := make([]promptcompiler.ExtensionCatalogGuidanceAction, 0, len(facts))
	for _, fact := range facts {
		actions = append(actions, promptcompiler.ExtensionCatalogGuidanceAction{ExtensionName: fact.ExtensionName, ActionName: fact.ActionName})
	}
	compiled, err := promptcompiler.CompileThreadContract(
		context.Background(),
		promptcompiler.EventThreadStart,
		"extension-e2e-thread-start",
		promptcompiler.CompileInput{
			AgentID: e2eAgent,
			Runtime: "codex",
			Extension: promptcompiler.ExtensionCompileInput{
				Enabled:  true,
				Guidance: promptcompiler.ExtensionCatalogGuidanceTemplateData{Actions: actions},
			},
		},
	)
	require.NoError(t, err)
	var promptBlock string
	for _, block := range compiled.Bundle.Blocks {
		if block.ID == "platform.agentteam.extension" {
			promptBlock = block.RenderedText()
			break
		}
	}
	require.NotEmpty(t, promptBlock)
	for _, secret := range appHostSecrets() {
		require.NotContains(t, string(serialized), secret)
		require.NotContains(t, help, secret)
		require.NotContains(t, guidance, secret)
		require.NotContains(t, promptBlock, secret)
	}
	require.Contains(t, promptBlock, "e2e-provider create-item")
}

type serverInvokeRecorder struct {
	mu               sync.Mutex
	requests         [][]byte
	catalogExchanges []catalogExchange
	next             http.Handler
}

type catalogExchange struct {
	Path        string
	IfNoneMatch string
	StatusCode  int
	ETag        string
	Body        []byte
}

func (recorder *serverInvokeRecorder) ServeHTTP(
	writer http.ResponseWriter,
	request *http.Request,
) {
	if request.Method == http.MethodGet &&
		strings.Contains(request.URL.EscapedPath(), "/runtime/ext") {
		response := httptest.NewRecorder()
		recorder.next.ServeHTTP(response, request)
		body := append([]byte(nil), response.Body.Bytes()...)
		recorder.mu.Lock()
		recorder.catalogExchanges = append(
			recorder.catalogExchanges,
			catalogExchange{
				Path:        request.URL.EscapedPath(),
				IfNoneMatch: request.Header.Get("If-None-Match"),
				StatusCode:  response.Code,
				ETag:        response.Header().Get("ETag"),
				Body:        body,
			},
		)
		recorder.mu.Unlock()
		for key, values := range response.Header() {
			for _, value := range values {
				writer.Header().Add(key, value)
			}
		}
		writer.WriteHeader(response.Code)
		_, _ = writer.Write(body)
		return
	}
	if request.Method == http.MethodPost &&
		strings.Contains(request.URL.EscapedPath(), "/runtime/ext/") {
		raw, err := io.ReadAll(request.Body)
		if err != nil {
			http.Error(writer, "read request", http.StatusInternalServerError)
			return
		}
		request.Body = io.NopCloser(bytes.NewReader(raw))
		recorder.mu.Lock()
		recorder.requests = append(recorder.requests, append([]byte(nil), raw...))
		recorder.mu.Unlock()
	}
	recorder.next.ServeHTTP(writer, request)
}

func (recorder *serverInvokeRecorder) catalogSnapshot() []catalogExchange {
	recorder.mu.Lock()
	defer recorder.mu.Unlock()
	out := make([]catalogExchange, len(recorder.catalogExchanges))
	for index, exchange := range recorder.catalogExchanges {
		out[index] = exchange
		out[index].Body = append([]byte(nil), exchange.Body...)
	}
	return out
}

func (recorder *serverInvokeRecorder) snapshot() [][]byte {
	recorder.mu.Lock()
	defer recorder.mu.Unlock()
	out := make([][]byte, len(recorder.requests))
	for index := range recorder.requests {
		out[index] = append([]byte(nil), recorder.requests[index]...)
	}
	return out
}

type catalogWireDocument struct {
	Extensions []catalogWireExtension `json:"extensions"`
}

type catalogWireExtension struct {
	ExtensionName string              `json:"extensionName"`
	DisplayName   string              `json:"displayName"`
	Actions       []catalogWireAction `json:"actions"`
}

type catalogWireAction struct {
	Name         string            `json:"name"`
	Summary      string            `json:"summary"`
	Description  string            `json:"description"`
	Safety       catalogWireSafety `json:"safety"`
	InputSchema  map[string]any    `json:"inputSchema"`
	OutputSchema map[string]any    `json:"outputSchema"`
	Example      map[string]any    `json:"example"`
}

type catalogWireSafety struct {
	Class string `json:"class"`
}

func requireCatalogResponse(
	t *testing.T,
	exchange catalogExchange,
	wantIfNoneMatch string,
	wantStatus int,
	wantETag string,
	wantExtensionNames []string,
) {
	t.Helper()
	require.True(t, strings.HasSuffix(exchange.Path, "/runtime/ext"))
	require.Equal(t, wantIfNoneMatch, exchange.IfNoneMatch)
	require.Equal(t, wantStatus, exchange.StatusCode)
	require.Equal(t, wantETag, exchange.ETag)
	require.NotEmpty(t, exchange.ETag)
	if wantStatus == http.StatusNotModified {
		require.Empty(t, exchange.Body)
		return
	}
	require.Equal(t, http.StatusOK, wantStatus)
	decoder := json.NewDecoder(bytes.NewReader(exchange.Body))
	decoder.DisallowUnknownFields()
	var document catalogWireDocument
	require.NoError(t, decoder.Decode(&document))
	require.Equal(t, io.EOF, decoder.Decode(&struct{}{}))
	names := make([]string, 0, len(document.Extensions))
	for _, extension := range document.Extensions {
		names = append(names, extension.ExtensionName)
		require.NotEmpty(t, extension.DisplayName)
		require.NotEmpty(t, extension.Actions)
		for _, action := range extension.Actions {
			require.NotEmpty(t, action.Name)
			require.NotEmpty(t, action.Summary)
			require.NotEmpty(t, action.Description)
			require.Equal(t, "mutation", action.Safety.Class)
			require.NotNil(t, action.InputSchema)
			require.NotNil(t, action.OutputSchema)
			require.NotNil(t, action.Example)
		}
	}
	require.ElementsMatch(t, wantExtensionNames, names)
}

func requireInvocationPayloadsHideAppHost(
	t *testing.T,
	requests [][]byte,
) {
	t.Helper()
	require.NotEmpty(t, requests)
	for _, raw := range requests {
		for _, secret := range appHostSecrets() {
			require.NotContains(t, string(raw), secret)
		}
		var payload struct {
			Input json.RawMessage `json:"input"`
		}
		decoder := json.NewDecoder(bytes.NewReader(raw))
		decoder.DisallowUnknownFields()
		require.NoError(t, decoder.Decode(&payload))
		require.NotEmpty(t, payload.Input)
	}
}

func appHostSecrets() []string {
	return []string{
		e2eHostDefinitionID,
		e2eHostInstallationID,
		e2eHostConnectionID,
		e2eHostSecondConnectionID,
		e2eHostSelectionID,
		e2eHostAuthorization,
		e2eHostSubject,
		e2eHostSecondSubject,
		"provider-secret",
		"provider-secret-second",
	}
}