package extensionflow

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/core/identifiers"
	"github.com/tidbcloud/AgentTeam/internal/extension/binding"
	"github.com/tidbcloud/AgentTeam/internal/extension/egress"
	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/store/extension"
)

// InvokeParams is the authoritative server invocation envelope. Workspace,
// Agent, and Agent member identity come from the authenticated runtime
// principal. Definition identity is resolved server-side from the durable
// catalog; Desktop never supplies the host App identity.
type InvokeParams struct {
	WorkspaceID   string
	AgentID       string
	AgentMemberID string
	ExtensionName string
	Action        string
	Input         json.RawMessage
}

// InvokeResult is the bounded, strategy-neutral, secret-free Agent result.
// Domain is the selected Agent-visible Extension name, not the provider's
// internal Definition domain. The result deliberately omits App identity,
// provider headers, URLs, credentials, connection identifiers, HTTP status,
// and raw response bytes.
type InvokeResult struct {
	Domain string          `json:"domain"`
	Action string          `json:"action"`
	Data   json.RawMessage `json:"data"`
}

type preparedInvocation struct {
	bound              binding.BoundRequest
	validatedInputJSON json.RawMessage
	requestHash        string
	requestBytes       int64
}

type invocationExecution struct {
	Data                 json.RawMessage
	AuditFailureData     json.RawMessage
	ExecutionStatusClass string
	ResultClassification string
	RequestBytes         int64
	ResponseBytes        int64
}

const defaultInvocationAuditTimeout = 5 * time.Second

func (service *Service) auditWriteTimeout() time.Duration {
	if service != nil && service.invocationAuditTimeout > 0 {
		return service.invocationAuditTimeout
	}
	return defaultInvocationAuditTimeout
}

// Invoke resolves every call back to its authorized App host, validates the
// pinned Action input, dispatches Builtin or External execution through one
// result budget, and persists one mandatory secret-free audit record.
func (service *Service) Invoke(
	ctx context.Context,
	params InvokeParams,
) (InvokeResult, error) {
	if err := service.validate(); err != nil {
		return InvokeResult{}, invocationInternal(
			"extension invocation is unavailable",
			err,
		)
	}
	if err := service.requireRegistry(); err != nil {
		return InvokeResult{}, invocationInternal(
			"extension invocation is unavailable",
			err,
		)
	}
	if err := validateInvokeParams(params); err != nil {
		return InvokeResult{}, err
	}
	if err := service.assertCredentialKey(); err != nil {
		return InvokeResult{}, invocationInternal(
			"extension invocation audit key is unavailable",
			err,
		)
	}

	startedAt := service.now().UTC()
	resolution, err := service.resolveInvocation(ctx, params)
	if err != nil {
		return InvokeResult{}, normalizeInvocationError(err)
	}
	prepared, err := prepareInvocation(resolution, params.Input)
	if err != nil {
		return InvokeResult{}, invocationError(
			InvocationValidation,
			"extension input failed validation",
			err,
		)
	}
	auditRequestHash, err := service.auditRequestHash(prepared.requestHash)
	if err != nil {
		return InvokeResult{}, invocationInternal(
			"extension invocation audit hash could not be created",
			err,
		)
	}

	execution, executionErr := service.executeResolvedInvocation(
		ctx,
		params,
		resolution,
		prepared,
		auditRequestHash,
	)
	result, classification, executionErr := finalizeInvocationExecution(
		resolution,
		&execution,
		executionErr,
	)

	auditCtx, cancelAudit := context.WithTimeout(
		context.WithoutCancel(ctx),
		service.auditWriteTimeout(),
	)
	defer cancelAudit()
	if err := service.writeInvocationAudit(
		auditCtx,
		params,
		resolution,
		prepared,
		auditRequestHash,
		execution,
		classification,
		startedAt,
	); err != nil {
		// Execution may already have applied a side effect. A trusted Builtin
		// may provide a conservative unknown result for exactly this
		// post-mutation failure window. Returning that result prevents a
		// generic 500 from falsely implying that retrying the mutation is safe.
		if executionErr == nil && len(execution.AuditFailureData) > 0 {
			fallback := execution
			fallback.Data = append(
				json.RawMessage(nil),
				execution.AuditFailureData...,
			)
			fallback.ResponseBytes = int64(len(fallback.Data))
			fallbackResult, _, fallbackErr := finalizeInvocationExecution(
				resolution,
				&fallback,
				nil,
			)
			if fallbackErr == nil {
				return fallbackResult, nil
			}
		}
		// Never return an apparent success without the required durable audit
		// record when no safe unknown projection is available.
		return InvokeResult{}, invocationInternal(
			"extension invocation outcome could not be recorded",
			err,
		)
	}
	if executionErr != nil {
		return InvokeResult{}, executionErr
	}
	return result, nil
}

func finalizeInvocationExecution(
	resolution invocationResolution,
	execution *invocationExecution,
	executionErr error,
) (InvokeResult, InvocationErrorCode, error) {
	result := InvokeResult{}
	if executionErr == nil {
		result = InvokeResult{
			Domain: resolution.selection.ExtensionName,
			Action: resolution.action.Name,
			Data:   append(json.RawMessage(nil), execution.Data...),
		}
		encoded, err := json.Marshal(result)
		switch {
		case err != nil:
			executionErr = invocationInternal(
				"extension result could not be encoded",
				err,
			)
		case len(encoded) > egress.MaxFinalResponseBytes:
			executionErr = invocationError(
				InvocationUpstreamFailure,
				"extension result exceeded the response budget",
				egress.ErrFinalResponseTooLarge,
			)
		}
	}
	if executionErr == nil {
		return result, "", nil
	}
	classification, _ := InvocationErrorCodeOf(executionErr)
	if classification == "" {
		classification = InvocationInternal
		executionErr = invocationInternal(
			"extension invocation failed",
			executionErr,
		)
	}
	if execution.ExecutionStatusClass == "" {
		execution.ExecutionStatusClass = string(classification)
	}
	return result, classification, executionErr
}

func validateInvokeParams(params InvokeParams) error {
	if params.WorkspaceID == "" ||
		strings.TrimSpace(params.WorkspaceID) != params.WorkspaceID ||
		params.AgentID == "" ||
		strings.TrimSpace(params.AgentID) != params.AgentID ||
		params.AgentMemberID == "" ||
		strings.TrimSpace(params.AgentMemberID) != params.AgentMemberID {
		return invocationError(
			InvocationPermissionDenied,
			"extension invocation principal is invalid",
			nil,
		)
	}
	if params.ExtensionName == "" ||
		strings.TrimSpace(params.ExtensionName) != params.ExtensionName ||
		params.Action == "" ||
		strings.TrimSpace(params.Action) != params.Action ||
		len(params.Input) == 0 {
		return invocationError(
			InvocationValidation,
			"extension invocation envelope is invalid",
			nil,
		)
	}
	if len(params.Input) > binding.MaxInputBytes {
		return invocationError(
			InvocationValidation,
			"extension input exceeded 1 MiB",
			binding.ErrInputTooLarge,
		)
	}
	return nil
}

func prepareInvocation(
	resolution invocationResolution,
	input json.RawMessage,
) (preparedInvocation, error) {
	switch resolution.kind {
	case registry.External:
		bound, err := binding.Bind(resolution.plan, input)
		if err != nil {
			return preparedInvocation{}, err
		}
		return preparedInvocation{
			bound:        bound,
			requestHash:  bound.SecretFreeHash,
			requestBytes: int64(len(bound.Body)),
		}, nil
	case registry.Builtin:
		validated, err := binding.ValidateActionInput(
			resolution.action.InputSchema,
			input,
		)
		if err != nil {
			return preparedInvocation{}, err
		}
		return preparedInvocation{
			validatedInputJSON: validated.ValidatedJSON,
			requestHash:        validated.SecretFreeHash,
			requestBytes:       int64(len(validated.ValidatedJSON)),
		}, nil
	default:
		return preparedInvocation{}, fmt.Errorf(
			"unsupported Extension execution kind %s",
			resolution.kind,
		)
	}
}

func (service *Service) executeResolvedInvocation(
	ctx context.Context,
	params InvokeParams,
	resolution invocationResolution,
	prepared preparedInvocation,
	auditRequestHash string,
) (invocationExecution, error) {
	switch resolution.kind {
	case registry.External:
		result, err := service.executeResolvedExternal(
			ctx,
			resolution,
			prepared.bound,
		)
		return invocationExecution{
			Data:                 append(json.RawMessage(nil), result.Data...),
			ExecutionStatusClass: result.StatusClass,
			RequestBytes:         result.RequestBytes,
			ResponseBytes:        result.ResponseBytes,
		}, err
	case registry.Builtin:
		executor, _, ok := service.registry.BuiltinExecutorFor(
			resolution.definition.Definition.Digest,
		)
		if !ok {
			return invocationExecution{
					RequestBytes: prepared.requestBytes,
				}, invocationInternal(
					"extension Builtin executor is unavailable",
					nil,
				)
		}
		result, err := executor.Invoke(ctx, registry.BuiltinInvocation{
			WorkspaceID:                    params.WorkspaceID,
			AgentID:                        params.AgentID,
			AgentMemberID:                  params.AgentMemberID,
			ServiceInstallationID:          resolution.selection.ServiceInstallationID,
			ServiceConnectionID:            resolution.connection.ID,
			ServiceConnectionGeneration:    resolution.connectionGeneration,
			ServiceInstallationExtensionID: resolution.selection.ID,
			AuthorizedByMemberID:           resolution.authorizedByMemberID,
			ExtensionName:                  resolution.selection.ExtensionName,
			Action:                         resolution.action.Name,
			DefinitionVersion: resolution.definition.Definition.
				Version,
			DefinitionDigest: resolution.definition.Definition.
				Digest,
			CatalogVersion: resolution.catalogVersion,
			RequestHash:    auditRequestHash,
			Input: append(
				json.RawMessage(nil),
				prepared.validatedInputJSON...,
			),
		})
		execution := invocationExecution{
			Data: append(json.RawMessage(nil), result.Data...),
			AuditFailureData: append(
				json.RawMessage(nil),
				result.AuditFailureData...,
			),
			ExecutionStatusClass: result.StatusClass,
			ResultClassification: result.ResultClassification,
			RequestBytes:         prepared.requestBytes,
			ResponseBytes:        int64(len(result.Data)),
		}
		if err != nil {
			execution.ExecutionStatusClass = "failed"
			if _, classified := InvocationErrorCodeOf(err); classified {
				return execution, err
			}
			if classified := classifyBuiltinExecutorError(err); classified != nil {
				return execution, classified
			}
			return execution, invocationError(
				InvocationUpstreamFailure,
				"extension Builtin execution failed",
				err,
			)
		}
		if execution.ExecutionStatusClass == "" {
			execution.ExecutionStatusClass = "success"
		}
		return execution, nil
	default:
		return invocationExecution{}, invocationInternal(
			"extension execution kind is invalid",
			nil,
		)
	}
}

type builtinInvocationError interface {
	error
	ExtensionInvocationErrorCode() string
}

func classifyBuiltinExecutorError(err error) error {
	var coded builtinInvocationError
	if !errors.As(err, &coded) || coded == nil {
		return nil
	}
	code := InvocationErrorCode(coded.ExtensionInvocationErrorCode())
	message := ""
	switch code {
	case InvocationValidation:
		message = "extension Builtin request is invalid"
	case InvocationPermissionDenied:
		message = "extension Builtin permission denied"
	case InvocationNotFound:
		message = "extension Builtin resource is unavailable"
	case InvocationConflict:
		message = "extension Builtin state conflicts with the request"
	case InvocationProviderAuthUnavailable:
		message = "extension Builtin provider authorization is unavailable"
	case InvocationRateLimited:
		message = "extension Builtin provider rate limited the request"
	case InvocationUpstreamFailure:
		message = "extension Builtin provider request failed"
	case InvocationInternal:
		message = "extension Builtin execution failed"
	default:
		return nil
	}
	return &InvocationError{Code: code, Message: message, cause: err}
}

const auditRequestHashDomain = "agentteam.extension.audit-request/v1\x00"

func (service *Service) auditRequestHash(
	secretFreeHash string,
) (string, error) {
	if err := service.assertCredentialKey(); err != nil {
		return "", err
	}
	if len(secretFreeHash) != len("sha256:")+sha256.Size*2 ||
		!strings.HasPrefix(secretFreeHash, "sha256:") {
		return "", errors.New("secret-free invocation hash is invalid")
	}
	mac := hmac.New(sha256.New, service.credentialKey)
	_, _ = mac.Write([]byte(auditRequestHashDomain))
	_, _ = mac.Write([]byte(secretFreeHash))
	return "sha256:" + hex.EncodeToString(mac.Sum(nil)), nil
}

func (service *Service) writeInvocationAudit(
	ctx context.Context,
	params InvokeParams,
	resolution invocationResolution,
	prepared preparedInvocation,
	requestHash string,
	result invocationExecution,
	classification InvocationErrorCode,
	startedAt time.Time,
) error {
	auditID, err := identifiers.NewEntityID("extaudit")
	if err != nil {
		return fmt.Errorf("generate extension invocation audit id: %w", err)
	}
	resultClassification := string(classification)
	if resultClassification == "" {
		resultClassification = result.ResultClassification
	}
	if resultClassification == "" {
		resultClassification = "success"
	}
	statusClass := result.ExecutionStatusClass
	if statusClass == "" {
		statusClass = "success"
	}
	requestBytes := result.RequestBytes
	if requestBytes == 0 {
		requestBytes = prepared.requestBytes
	}
	record := extension.InvocationAudit{
		ID:                             auditID,
		WorkspaceID:                    params.WorkspaceID,
		AgentID:                        params.AgentID,
		ServiceInstallationExtensionID: resolution.selection.ID,
		ExtensionName:                  resolution.selection.ExtensionName,
		Domain:                         resolution.definition.Definition.Domain,
		Action:                         resolution.action.Name,
		DefinitionVersion:              resolution.definition.Definition.Version,
		DefinitionDigest:               resolution.definition.Definition.Digest,
		CatalogVersion:                 resolution.catalogVersion,
		RequestHash:                    requestHash,
		ExecutionStatusClass:           statusClass,
		ResultClassification:           resultClassification,
		Latency:                        service.now().UTC().Sub(startedAt),
		RequestBytes:                   requestBytes,
		ResponseBytes:                  result.ResponseBytes,
		OccurredAt:                     startedAt,
	}
	err = service.store.WriteInvocationAuditOnce(ctx, record)
	if err == nil || ctx.Err() != nil {
		return err
	}
	// The first insert may have committed even if its acknowledgement was
	// lost. Retry only the audit with the same caller-owned ID; first-fact-wins
	// makes this safe and never replays provider execution.
	return service.store.WriteInvocationAuditOnce(ctx, record)
}

func normalizeInvocationError(err error) error {
	if err == nil {
		return nil
	}
	if _, ok := InvocationErrorCodeOf(err); ok {
		return err
	}
	return invocationInternal("extension invocation failed", err)
}