package exec

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

func handleRuntimeCatalog(adapter adapter) gin.HandlerFunc {
	return func(c *gin.Context) {
		authorized, release, ok := acquireRuntimeRequest(c, adapter, "Catalog")
		if !ok {
			return
		}
		defer release()
		view, err := authorized.RuntimeCatalog.View()
		if err != nil {
			writeError(c, ApplyErrorBudget(&Error{Code: CodeInternal, Problem: "internal error", Cause: err}))
			return
		}
		raw, err := json.Marshal(runtimeCatalogDocumentFromView(view))
		if err != nil || len(raw) > maxRuntimeCatalogBytes {
			writeError(c, ApplyErrorBudget(NewError(CodeInternal, "internal error")))
			return
		}
		if execAuthorizationRevoked(authorized.revoked) {
			writeError(c, ApplyErrorBudget(NewError(CodePermissionDenied, "Exec Catalog scope mismatch")))
			return
		}
		c.Data(http.StatusOK, "application/json", raw)
	}
}

func handleRuntimeAction(adapter adapter) gin.HandlerFunc {
	return func(c *gin.Context) {
		handleActionRequest(
			c,
			adapter,
			strings.TrimSpace(c.Param("domain")),
			strings.TrimSpace(c.Param("action")),
		)
	}
}

func handleActionRequest(c *gin.Context, adapter adapter, domain, actionName string) {
	authorized, release, ok := acquireRuntimeRequest(c, adapter, "action")
	if !ok {
		return
	}
	defer release()
	action, found := authorized.RuntimeCatalog.Lookup(domain, actionName)
	if !found {
		writeError(c, ApplyErrorBudget(NewError(CodeNotFound, "unknown action")))
		return
	}
	requestContext, stopWatch := watchRuntimeRevocation(c, authorized.revoked)
	defer stopWatch()
	raw, err := readAuthorizedActionBody(c, adapter, authorized.revoked)
	if err != nil {
		if execAuthorizationRevoked(authorized.revoked) {
			writeError(c, ApplyErrorBudget(NewError(CodePermissionDenied, "Exec action scope mismatch")))
		} else {
			writeError(c, ApplyErrorBudget(err))
		}
		return
	}
	if err := validateReplayHeaders(domain, actionName, raw, c.GetHeader("Idempotency-Key"), c.GetHeader("X-Request-ID")); err != nil {
		writeError(c, ApplyErrorBudget(err))
		return
	}
	var encoded []byte
	switch action.Source {
	case RuntimeActionBuiltin:
		invoker, exists := adapter.builtin[domain+" "+actionName]
		if !exists || invoker.invoke == nil {
			err = NewError(CodeInternal, "internal error")
		} else {
			encoded, err = invoker.call(requestContext, authorized, raw)
		}
	case RuntimeActionExtension:
		encoded, err = invokeExtensionAction(
			requestContext,
			authorized,
			action,
			raw,
		)
	default:
		err = NewError(CodeInternal, "internal error")
	}
	if err != nil {
		if execAuthorizationRevoked(authorized.revoked) || errors.Is(err, context.Canceled) {
			writeError(c, ApplyErrorBudget(NewError(CodePermissionDenied, "Exec action scope mismatch")))
		} else {
			writeError(c, ApplyErrorBudget(err))
		}
		return
	}
	if execAuthorizationRevoked(authorized.revoked) {
		writeError(c, ApplyErrorBudget(NewError(CodePermissionDenied, "Exec action scope mismatch")))
		return
	}
	c.Data(http.StatusOK, "application/json", encoded)
}

func acquireRuntimeRequest(
	c *gin.Context,
	adapter adapter,
	kind string,
) (runtimeAuthorization, func(), bool) {
	if adapter.authorizer == nil {
		writeError(c, ApplyErrorBudget(NewError(CodeInternal, "internal error")))
		return runtimeAuthorization{}, nil, false
	}
	tokenValue, err := strictBearer(c.GetHeader("Authorization"))
	if err != nil {
		writeError(c, ApplyErrorBudget(err))
		return runtimeAuthorization{}, nil, false
	}
	serverID := strings.TrimSpace(c.Param("serverID"))
	agentID := strings.TrimSpace(c.Param("agentID"))
	if serverID == "" || agentID == "" {
		writeError(c, ApplyErrorBudget(NewError(CodePermissionDenied, "Exec "+kind+" scope mismatch")))
		return runtimeAuthorization{}, nil, false
	}
	authorized, release, err := adapter.authorizer.acquire(c.Request.Context(), tokenValue, serverID, agentID)
	if err == nil {
		return authorized, release, true
	}
	if errors.Is(err, errExecAuthorizationCapacity) {
		writeError(c, ApplyErrorBudget(NewError(CodeConflict, "Exec runtime authorization is busy")))
		return runtimeAuthorization{}, nil, false
	}
	logExecAuthorizationFailure(err)
	if isKnownExecAuthorizationFailure(err) {
		writeError(c, ApplyErrorBudget(&Error{
			Code: CodePermissionDenied, Problem: "Exec " + kind + " scope mismatch", Cause: err,
		}))
	} else {
		writeError(c, ApplyErrorBudget(&Error{Code: CodeInternal, Problem: "internal error", Cause: err}))
	}
	return runtimeAuthorization{}, nil, false
}

func watchRuntimeRevocation(
	c *gin.Context,
	revoked <-chan struct{},
) (context.Context, func()) {
	requestContext, cancelRequest := context.WithCancel(c.Request.Context())
	stop := make(chan struct{})
	done := make(chan struct{})
	go func() {
		defer close(done)
		select {
		case <-revoked:
			cancelRequest()
			_ = c.Request.Body.Close()
		case <-stop:
		}
	}()
	return requestContext, func() {
		close(stop)
		<-done
		cancelRequest()
	}
}

type extensionInvokeRequest struct {
	Input json.RawMessage `json:"input"`
}

type extensionInvokeResult struct {
	Domain string          `json:"domain"`
	Action string          `json:"action"`
	Data   json.RawMessage `json:"data"`
}

func invokeExtensionAction(
	ctx context.Context,
	authorized runtimeAuthorization,
	runtimeAction RuntimeAction,
	raw json.RawMessage,
) ([]byte, error) {
	descriptor, ok := runtimeAction.Extension()
	if !ok {
		return nil, NewError(CodeInternal, "internal error")
	}
	if err := descriptor.ValidateInput(raw); err != nil {
		return nil, &Error{Code: CodeValidation, Problem: "invalid input", Cause: err}
	}
	var response json.RawMessage
	err := authorized.Server.PostWithHeadersResponseLimit(
		ctx,
		authorized.Server.AgentRuntimePath(
			"/ext/"+url.PathEscape(runtimeAction.Domain)+
				"/actions/"+url.PathEscape(runtimeAction.Name),
		),
		extensionInvokeRequest{Input: raw},
		&response,
		nil,
		maxSerializedResultBytes,
	)
	if err != nil {
		return nil, classifyExtensionUpstreamError(err, runtimeAction.Domain, runtimeAction.Name)
	}
	result, err := decodeExtensionInvokeResult(response)
	if err != nil ||
		result.Domain != runtimeAction.Domain ||
		result.Action != runtimeAction.Name {
		return nil, &Error{
			Code: CodeInternal, Problem: "Extension invocation returned an invalid response",
			Cause: err,
		}
	}
	outputValidator, err := compileExtensionSchema(descriptor.OutputSchema)
	if err != nil {
		return nil, &Error{Code: CodeInternal, Problem: "internal error", Cause: err}
	}
	if outputValidator != nil {
		if err := validateInstance(outputValidator, result.Data); err != nil {
			return nil, &Error{
				Code: CodeInternal, Problem: "Extension invocation returned an invalid response",
				Cause: err,
			}
		}
	}
	encoded, err := json.Marshal(result)
	if err != nil {
		return nil, &Error{
			Code: CodeInternal, Problem: "Extension invocation returned an invalid response",
			Cause: err,
		}
	}
	if len(encoded) > maxSerializedResultBytes {
		return nil, resultBudgetExceededError(runtimeAction.Domain, runtimeAction.Name, len(encoded), nil)
	}
	return encoded, nil
}

func decodeExtensionInvokeResult(raw json.RawMessage) (extensionInvokeResult, error) {
	var result extensionInvokeResult
	decoder := json.NewDecoder(strings.NewReader(string(raw)))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&result); err != nil {
		return extensionInvokeResult{}, err
	}
	if err := decoder.Decode(new(struct{})); err != io.EOF {
		return extensionInvokeResult{}, errors.New("Extension invocation returned trailing JSON")
	}
	if strings.TrimSpace(result.Domain) == "" ||
		strings.TrimSpace(result.Action) == "" ||
		len(result.Data) == 0 ||
		!json.Valid(result.Data) {
		return extensionInvokeResult{}, errors.New("Extension invocation result is incomplete")
	}
	return result, nil
}

func classifyExtensionUpstreamError(err error, action ...string) error {
	if errors.Is(err, apiclient.ErrResponseTooLarge) {
		domain, name := "extension", "action"
		if len(action) >= 2 {
			domain, name = action[0], action[1]
		}
		return resultBudgetExceededError(domain, name, maxSerializedResultBytes+1, err)
	}
	var upstream *apiclient.APIError
	if !errors.As(err, &upstream) {
		return &Error{Code: CodeInternal, Problem: "Extension invocation failed", Cause: err}
	}
	// The agent runtime rejects a rotated-out daemon token with HTTP 401 and a
	// body that carries no machine-readable code, so classify by status. An
	// authentication refusal must not be reported as a generic internal
	// failure: it reads as an extension outage and hides the token-lifecycle
	// root cause.
	if upstream.StatusCode == http.StatusUnauthorized {
		return &Error{
			Code:    CodePermissionDenied,
			Problem: "Extension invocation failed (agent authentication rejected)",
			Cause:   err,
		}
	}
	code := CodeInternal
	switch upstream.Code {
	case "validation":
		code = CodeValidation
	case "permission_denied":
		code = CodePermissionDenied
	case "not_found":
		code = CodeNotFound
	case "conflict":
		code = CodeConflict
	case "agent_extension_catalog_stale":
		code = CodeAgentExtensionCatalogStale
	case "rate_limited":
		code = CodeRateLimited
	case "provider_auth_unavailable":
		code = CodeProviderAuthUnavailable
	case "upstream_failure":
		code = CodeUpstreamFailure
	case "internal":
		code = CodeInternal
	}
	return &Error{Code: code, Problem: "Extension invocation failed", Cause: err}
}