package extensionflow

import (
	"bytes"
	"context"
	"errors"
	"fmt"

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

// invocationResolution is the server-internal result of authoritative
// Extension resolution. Desktop supplies only the flattened Extension name;
// the host App selection, installation, connection, and execution kind are
// derived and verified here.
type invocationResolution struct {
	selection  serviceactor.InstallationExtension
	definition registry.RegisteredExtension
	action     registry.DefinitionAction
	kind       registry.Kind

	catalogVersion int64

	plan                 registry.ActionPlan
	connection           extension.RuntimeConnection
	connectionGeneration string
	hasConnection        bool
	encryptedCredentials []byte
	authorizedByMemberID string
}

func (service *Service) resolveInvocation(
	ctx context.Context,
	params InvokeParams,
) (invocationResolution, error) {
	var resolved invocationResolution
	err := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
		var err error
		resolved, err = service.resolveInvocationTx(
			ctx,
			tx,
			service.store.InTx(tx),
			params,
		)
		return err
	})
	if err != nil {
		return invocationResolution{}, err
	}
	return resolved, nil
}

func (service *Service) resolveInvocationTx(
	ctx context.Context,
	tx coredb.Tx,
	extTx *extension.TxStore,
	params InvokeParams,
) (invocationResolution, error) {
	// The catalog row is locked BEFORE the Agent membership assertion. The
	// catalog-update event INSERT takes a foreign-key shared lock on the
	// parent agents row; AssertWorkspaceAgentMember takes an exclusive lock
	// on that same row. Holding the membership lock while waiting for the
	// catalog row therefore deadlocks with a concurrent authorization change
	// (catalog writer holds the catalog row, waits for the agents-row FK
	// lock). Locking the catalog row first keeps this transaction lock-free
	// while it waits, so the catalog writer can always drain and commit.
	effectiveExtensions, catalogVersion, err := service.verifyInvocationCatalog(
		ctx,
		tx,
		extTx,
		params,
	)
	if err != nil {
		return invocationResolution{}, err
	}
	if err := extTx.AssertWorkspaceAgentMember(
		ctx,
		params.WorkspaceID,
		params.AgentID,
		params.AgentMemberID,
	); err != nil {
		return invocationResolution{}, invocationError(
			InvocationPermissionDenied,
			"extension invocation principal is not authorized",
			err,
		)
	}
	resolved, err := selectInvocationExtension(
		effectiveExtensions,
		params,
	)
	if err != nil {
		return invocationResolution{}, err
	}
	resolved.catalogVersion = catalogVersion
	resolved, err = resolveInvocationConnection(
		ctx,
		extTx,
		params.WorkspaceID,
		resolved,
		effectiveExtensions,
	)
	if err != nil {
		return invocationResolution{}, err
	}
	return service.resolveInvocationStrategy(resolved, params)
}

func (service *Service) verifyInvocationCatalog(
	ctx context.Context,
	tx coredb.Tx,
	extTx *extension.TxStore,
	params InvokeParams,
) ([]effectiveExtension, int64, error) {
	// Authorization decisions are read as a non-locking snapshot BEFORE the
	// catalog row is locked. Locking them FOR UPDATE would deadlock with a
	// concurrent revoke/disable: the principal authorization index carries
	// the status column, so a lifecycle UPDATE must replace index entries an
	// ordered locking scan holds. The catalog row is the invocation path's
	// single lock point, ordered after the delivery/reconcile paths'
	// authorization locks; the durable catalog-state compare below fences
	// staleness (a decision change always lands with a matching catalog
	// mutation in the same writer transaction, so a snapshot/durable
	// mismatch fails closed as stale instead of authorizing drift).
	decisions, err := extTx.ListAgentAppAuthorizationDecisionsSnapshot(
		ctx,
		params.WorkspaceID,
		params.AgentID,
	)
	if err != nil {
		return nil, 0, invocationInternal(
			"App-hosted extension authorization could not be verified",
			err,
		)
	}
	if service.decisionsSnapshotHook != nil {
		service.decisionsSnapshotHook()
	}
	catalog, found, err := extTx.LoadAgentExtensionCatalogForUpdate(
		ctx,
		params.WorkspaceID,
		params.AgentID,
	)
	if err != nil {
		return nil, 0, invocationInternal(
			"agent extension catalog could not be verified",
			err,
		)
	}
	if !found {
		return nil, 0, invocationError(
			InvocationAgentCatalogStale,
			"agent extension catalog is stale",
			nil,
		)
	}
	effectiveExtensions, err := service.resolveAppHostedExtensionsFromDecisions(
		ctx,
		tx,
		extTx,
		params.WorkspaceID,
		params.AgentID,
		decisions,
	)
	if err != nil {
		return nil, 0, invocationInternal(
			"App-hosted extension authorization could not be verified",
			err,
		)
	}
	entries, err := catalogEntriesForEffectiveExtensions(effectiveExtensions)
	if err != nil {
		return nil, 0, invocationInternal(
			"agent extension catalog could not be projected",
			err,
		)
	}
	recomputed, err := marshalCatalogState(entries)
	if err != nil {
		return nil, 0, invocationInternal(
			"agent extension catalog state could not be verified",
			err,
		)
	}
	if catalog.Enabled != (len(entries) > 0) ||
		catalog.State != recomputed {
		return nil, 0, invocationError(
			InvocationAgentCatalogStale,
			"agent extension catalog is stale",
			errors.New(
				"durable catalog state does not match the live App-hosted surface",
			),
		)
	}
	return effectiveExtensions, catalog.Version, nil
}

func selectInvocationExtension(
	effectiveExtensions []effectiveExtension,
	params InvokeParams,
) (invocationResolution, error) {
	effective, found, err := findEffectiveExtensionByName(
		effectiveExtensions,
		params.ExtensionName,
	)
	if err != nil {
		return invocationResolution{}, invocationInternal(
			"extension selection is ambiguous",
			err,
		)
	}
	if !found {
		return invocationResolution{}, invocationError(
			InvocationNotFound,
			"extension is not available",
			nil,
		)
	}
	action, found := findDefinitionAction(
		effective.Definition.Definition.Actions,
		params.Action,
	)
	if !found || !containsExactString(effective.Actions, params.Action) {
		return invocationResolution{}, invocationError(
			InvocationNotFound,
			"extension action is not available",
			nil,
		)
	}
	return invocationResolution{
		selection:            effective.Selection,
		definition:           effective.Definition,
		action:               action,
		kind:                 effective.Definition.Kind,
		authorizedByMemberID: effective.AuthorizedByMemberID,
	}, nil
}

func containsExactString(values []string, target string) bool {
	for _, value := range values {
		if value == target {
			return true
		}
	}
	return false
}

func resolveInvocationConnection(
	ctx context.Context,
	extTx *extension.TxStore,
	workspaceID string,
	resolved invocationResolution,
	effectiveExtensions []effectiveExtension,
) (invocationResolution, error) {
	effective, found, err := findEffectiveExtensionByName(
		effectiveExtensions,
		resolved.selection.ExtensionName,
	)
	if err != nil || !found {
		return invocationResolution{}, invocationInternal(
			"extension selection changed during resolution",
			err,
		)
	}
	switch resolved.definition.ConnectionPolicy {
	case registry.ConnectionNone:
		if len(effective.Connections) != 0 {
			return invocationResolution{}, invocationInternal(
				"connectionless extension has an unexpected provider connection",
				fmt.Errorf(
					"expected zero active App connections, got %d",
					len(effective.Connections),
				),
			)
		}
		return resolved, nil
	case registry.ConnectionRequired:
		// Continue through the existing exactly-one connection and credential
		// resolution path below.
	default:
		return invocationResolution{}, invocationInternal(
			"extension connection policy is invalid",
			fmt.Errorf(
				"unsupported connection policy %d",
				resolved.definition.ConnectionPolicy,
			),
		)
	}
	if len(effective.Connections) != 1 {
		return invocationResolution{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider authentication is unavailable",
			fmt.Errorf(
				"expected exactly one active App connection, got %d",
				len(effective.Connections),
			),
		)
	}
	connections, err := extTx.ListActiveAppRuntimeConnections(
		ctx,
		resolved.selection.ServiceInstallationID,
	)
	if err != nil {
		return invocationResolution{}, invocationInternal(
			"extension provider connection could not be loaded",
			err,
		)
	}
	if len(connections) != 1 {
		return invocationResolution{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider authentication is unavailable",
			fmt.Errorf(
				"expected exactly one active App runtime connection, got %d",
				len(connections),
			),
		)
	}
	connection := connections[0]
	if err := validateAppRuntimeConnection(
		connection,
		effective.Connections[0],
		workspaceID,
		resolved.selection,
	); err != nil {
		return invocationResolution{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider authentication is unavailable",
			err,
		)
	}
	resolved.connection = connection
	// The generation is derived from the CURRENT locked read, which the
	// field-by-field compare above just proved identical to the catalog
	// snapshot (subject, credential digest, and config included) — a
	// concurrent reconnect can never pair fresh credentials with a stale
	// generation.
	resolved.connectionGeneration = extension.RuntimeConnectionGeneration(
		connection.ID,
		connection.ExternalSubjectRef,
		connection.CurrentCredentialDigest,
		connection.Config,
	)
	resolved.hasConnection = true
	resolved.encryptedCredentials = append(
		[]byte(nil),
		connection.Credentials...,
	)
	return resolved, nil
}

func (service *Service) resolveInvocationStrategy(
	resolved invocationResolution,
	params InvokeParams,
) (invocationResolution, error) {
	switch resolved.kind {
	case registry.External:
		plan, ok := service.registry.ActionPlan(
			resolved.definition.Definition.Digest,
			params.Action,
		)
		if !ok ||
			plan.Domain != resolved.definition.Definition.Domain ||
			plan.DefinitionVersion != resolved.definition.Definition.Version ||
			plan.DefinitionDigest != resolved.definition.Definition.Digest {
			return invocationResolution{}, invocationInternal(
				"extension action plan is unavailable",
				errors.New(
					"digest-pinned External action plan does not match the selection",
				),
			)
		}
		resolved.plan = plan
	case registry.Builtin:
		if _, _, ok := service.registry.BuiltinExecutorFor(
			resolved.definition.Definition.Digest,
		); !ok {
			return invocationResolution{}, invocationInternal(
				"extension Builtin executor is unavailable",
				nil,
			)
		}
	default:
		return invocationResolution{}, invocationInternal(
			"extension execution kind is invalid",
			fmt.Errorf("unsupported registry kind %s", resolved.kind),
		)
	}
	return resolved, nil
}

func findEffectiveExtensionByName(
	extensions []effectiveExtension,
	extensionName string,
) (effectiveExtension, bool, error) {
	var selected effectiveExtension
	found := false
	for _, extension := range extensions {
		if extension.Selection.ExtensionName != extensionName {
			continue
		}
		if found {
			return effectiveExtension{}, false, fmt.Errorf(
				"multiple active selections use Extension name %q",
				extensionName,
			)
		}
		selected = extension
		found = true
	}
	return selected, found, nil
}

func findDefinitionAction(
	actions []registry.DefinitionAction,
	actionName string,
) (registry.DefinitionAction, bool) {
	for _, action := range actions {
		if action.Name == actionName {
			return action, true
		}
	}
	return registry.DefinitionAction{}, false
}

func validateAppRuntimeConnection(
	connection extension.RuntimeConnection,
	projected extension.AppHostedConnectionIdentity,
	workspaceID string,
	selection serviceactor.InstallationExtension,
) error {
	if connection.ID == "" ||
		connection.ID != projected.ID ||
		connection.WorkspaceID != workspaceID ||
		connection.WorkspaceID != projected.WorkspaceID ||
		connection.ServiceDefinitionID != selection.ServiceDefinitionID ||
		connection.ServiceDefinitionID != projected.ServiceDefinitionID ||
		connection.ServiceInstallationID != selection.ServiceInstallationID ||
		connection.ServiceInstallationID != projected.ServiceInstallationID ||
		connection.ServiceKind != serviceactor.KindApp ||
		connection.ServiceKind != projected.ServiceKind ||
		connection.ExternalSubjectRef == "" ||
		connection.Status != serviceactor.StatusActive ||
		connection.Status != projected.Status ||
		len(connection.Credentials) == 0 {
		return errors.New(
			"App runtime connection identity/status is inconsistent",
		)
	}
	// The locked current read must be the SAME generation the catalog
	// snapshot pinned: subject, credential digest, and config are compared
	// field-by-field so a concurrent reconnect that landed between the
	// snapshot and this lock fails closed instead of mixing generations.
	if connection.ExternalSubjectRef != projected.ExternalSubjectRef ||
		connection.CurrentCredentialDigest != projected.CurrentCredentialDigest ||
		!bytes.Equal(connection.Config, projected.Config) {
		return errors.New(
			"App runtime connection changed under the catalog snapshot",
		)
	}
	return nil
}