package extension

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/service/extensionflow"
)

// Handler serves the Agent-visible Extension Catalog and the invoke route.
type Handler struct {
	catalog  CatalogSource
	access   AuthorizationSource
	snapshot CatalogSnapshotSource
	invoker  InvocationService
}

// NewHandler returns a Handler that projects the Agent-visible Catalog from the
// immutable Registry, filtered by the Agent's authorized domains/actions (stage 6).
func NewHandler(dependencies Dependencies) *Handler {
	return &Handler{
		catalog: dependencies.Catalog, access: dependencies.Access,
		snapshot: dependencies.Snapshot,
		invoker:  dependencies.Invoker,
	}
}

// authorizedExtensions is the Agent's resolved Agent Access, shared by discovery
// fallback and invocation to avoid authorization drift. It is keyed by the
// explicit Agent-visible Extension name so multi-instance installations of one
// Definition do not collapse.
type authorizedExtensions struct {
	entries map[string]AgentDomainAccess
}

// loadAuthorizedExtensions resolves the requesting Agent's authorized Extension domains.
// Fail-closed: nil access or error → no visible extensions (no domains/actions visible).
func (handler *Handler) loadAuthorizedExtensions(request api.AgentScopedRequest) (*authorizedExtensions, error) {
	if handler.access == nil {
		return &authorizedExtensions{entries: map[string]AgentDomainAccess{}}, nil
	}
	principal := request.Principal
	authorized, err := handler.access.ActiveAgentDomains(request.Gin.Request.Context(), principal.WorkspaceID, principal.AgentID)
	if err != nil {
		return nil, err
	}
	surface := &authorizedExtensions{entries: make(map[string]AgentDomainAccess, len(authorized))}
	for _, access := range authorized {
		key := access.ExtensionName
		if key == "" {
			return nil, errors.New("extension access has no Extension name")
		}
		// Fail closed on a duplicate Extension name rather than silently letting
		// the second overwrite the first, which would widen or narrow the
		// Agent-visible action set.
		if _, exists := surface.entries[key]; exists {
			return nil, fmt.Errorf("duplicate Extension name %q for agent", key)
		}
		surface.entries[key] = access
	}
	return surface, nil
}

func (s *authorizedExtensions) hasEntry(key string) bool {
	_, ok := s.entries[key]
	return ok
}

func (s *authorizedExtensions) accessFor(key string) (AgentDomainAccess, bool) {
	access, ok := s.entries[key]
	return access, ok
}

func (s *authorizedExtensions) hasAction(key, action string) bool {
	access, ok := s.entries[key]
	if !ok {
		return false
	}
	for _, a := range access.Actions {
		if a == action {
			return true
		}
	}
	return false
}

// Root returns the complete authorized Catalog document (spec §6.2 root GET):
// every authorized Extension with the FULL descriptors of its authorized
// Actions. Fail-closed: nil access → empty catalog (NOT the full registry).
func (handler *Handler) Root(request api.AgentScopedRequest) (api.Response, error) {
	// Consistent snapshot path: entries verified against the durable catalog row
	// in one transaction so a concurrent Grant/Revoke cannot return a stale
	// entry set (stage-6 final-review blocker 4).
	if handler.snapshot != nil {
		snap, err := handler.snapshot.AgentCatalogSnapshot(request.Gin.Request.Context(), request.Principal.WorkspaceID, request.Principal.AgentID)
		if err != nil {
			return api.Response{}, api.NewHTTPError(http.StatusInternalServerError, "failed to read agent catalog snapshot")
		}
		return conditionalCatalogResponse(
			request.Gin.GetHeader("If-None-Match"),
			projectAuthorizedRoot(handler.catalog.List(), snap.Domains),
		)
	}
	// Fallback: separate loadAuthorizedExtensions (for test setups without Snapshot).
	surface, err := handler.loadAuthorizedExtensions(request)
	if err != nil {
		return api.Response{}, api.NewHTTPError(http.StatusInternalServerError, "failed to load agent authorizations")
	}
	authorized := make([]AgentDomainAccess, 0, len(surface.entries))
	for _, access := range surface.entries {
		authorized = append(authorized, access)
	}
	return conditionalCatalogResponse(
		request.Gin.GetHeader("If-None-Match"),
		projectAuthorizedRoot(handler.catalog.List(), authorized),
	)
}

func conditionalCatalogResponse(
	ifNoneMatch string,
	root rootResponse,
) (api.Response, error) {
	encoded, err := json.Marshal(root)
	if err != nil {
		return api.Response{}, api.NewHTTPErrorWithCause(
			http.StatusInternalServerError,
			"failed to encode agent extension catalog",
			err,
		)
	}
	sum := sha256.Sum256(encoded)
	etag := `"ext-v1-` +
		base64.RawURLEncoding.EncodeToString(sum[:]) +
		`"`
	headers := make(http.Header)
	// no-transform keeps the strong ETag bound to the exact JSON bytes all the
	// way through the shared compression middleware and intermediaries.
	headers.Set("Cache-Control", "private, no-cache, no-transform")
	headers.Set("ETag", etag)
	headers.Set("Vary", "Authorization")
	if weakETagMatch(ifNoneMatch, etag) {
		return api.Response{
			Status:  http.StatusNotModified,
			NoBody:  true,
			Headers: headers,
		}, nil
	}
	return api.Response{
		Status:  http.StatusOK,
		Body:    json.RawMessage(encoded),
		Headers: headers,
	}, nil
}

// weakETagMatch implements the weak comparison required by If-None-Match for
// GET. The server emits a strong content-derived ETag, but a conforming client
// may send the equivalent W/ form or a comma-separated validator list.
func weakETagMatch(ifNoneMatch string, current string) bool {
	current = strings.TrimPrefix(strings.TrimSpace(current), "W/")
	if current == "" {
		return false
	}
	for _, candidate := range strings.Split(ifNoneMatch, ",") {
		candidate = strings.TrimSpace(candidate)
		if candidate == "*" {
			return true
		}
		if strings.TrimPrefix(candidate, "W/") == current {
			return true
		}
	}
	return false
}

// invokeRequest carries only Action input. No catalog or definition identity
// crosses the wire: the server resolves every durable execution invariant
// itself. Input is a json.RawMessage so any valid JSON value is accepted
// losslessly; schema validation happens in the execution stage.
type invokeRequest struct {
	Input json.RawMessage `json:"input"`
}

// invokeEnvelopeCapBytes is the total invoke-envelope byte cap enforced AFTER
// authentication and BEFORE decode (spec §6.4 transport boundary). It carries
// headroom above the 1 MiB Action input cap so a legitimate 1 MiB input plus
// its envelope overhead does not trip the transport threshold; per-Action
// input size/schema validation happens in the execution stage.
const invokeEnvelopeCapBytes = 2 << 20

// Invoke authenticates and bounds the transport envelope, hides unauthorized
// path identities, then delegates to the authoritative service. The service
// rechecks every durable execution invariant; this discovery precheck is not an
// authorization substitute.
func (handler *Handler) Invoke(request api.AgentScopedRequest) (api.Response, error) {
	key := request.Gin.Param("extensionName")
	actionName := request.Gin.Param("action")
	surface, err := handler.loadAuthorizedExtensions(request)
	if err != nil {
		return api.Response{}, api.NewHTTPError(http.StatusInternalServerError, "failed to load agent authorizations")
	}
	access, ok := surface.accessFor(key)
	if !ok {
		return api.Response{}, api.NewHTTPErrorWithCode(http.StatusNotFound, "extension domain not found", extensionNotFoundCode)
	}
	if !surface.hasAction(key, actionName) {
		return api.Response{}, api.NewHTTPErrorWithCode(http.StatusNotFound, "extension action not found", extensionNotFoundCode)
	}
	ext, ok := definitionByIdentity(
		handler.catalog.List(),
		access.Domain,
		access.DefinitionVersion,
		access.DefinitionDigest,
	)
	if !ok {
		return api.Response{}, api.NewHTTPErrorWithCode(http.StatusNotFound, "extension domain not found", extensionNotFoundCode)
	}
	if _, ok := findAction(ext, actionName); !ok {
		return api.Response{}, api.NewHTTPErrorWithCode(http.StatusNotFound, "extension action not found", extensionNotFoundCode)
	}
	data, err := readInvokeEnvelope(request.Gin.Request.Body)
	if err != nil {
		return api.Response{}, err
	}
	envelope, err := decodeInvokeEnvelope(data)
	if err != nil {
		return api.Response{}, err
	}
	if handler.invoker == nil {
		return api.Response{}, api.NewHTTPErrorWithCode(
			http.StatusInternalServerError,
			"extension invocation is unavailable",
			string(extensionflow.InvocationInternal),
		)
	}
	result, err := handler.invoker.Invoke(
		request.Gin.Request.Context(),
		extensionflow.InvokeParams{
			WorkspaceID:   request.Principal.WorkspaceID,
			AgentID:       request.Principal.AgentID,
			AgentMemberID: request.Principal.MemberID,
			ExtensionName: key,
			Action:        actionName,
			Input:         append(json.RawMessage(nil), envelope.Input...),
		},
	)
	if err != nil {
		return api.Response{}, classifyInvocationError(err)
	}
	return api.OK(result), nil
}

// readInvokeEnvelope reads at most invokeEnvelopeCapBytes of the authenticated
// request body, BEFORE decode (spec §6.4). An over-cap body is rejected as 413
// Payload Too Large rather than decoded: the body cap is a transport boundary,
// not a content check. It returns the capped bytes ready for single-document
// decoding.
func readInvokeEnvelope(body io.Reader) ([]byte, error) {
	if body == nil {
		return nil, api.NewHTTPErrorWithCode(http.StatusBadRequest, "invalid request body", extensionInvalidBodyCode)
	}
	reader := io.LimitReader(body, invokeEnvelopeCapBytes+1)
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, api.NewHTTPErrorWithCode(http.StatusBadRequest, "invalid request body", extensionInvalidBodyCode)
	}
	if len(data) > invokeEnvelopeCapBytes {
		return nil, api.NewHTTPErrorWithCode(http.StatusRequestEntityTooLarge, "request body too large", extensionBodyTooLargeCode)
	}
	return data, nil
}

// decodeInvokeEnvelope decodes EXACTLY one JSON document from data into the
// invoke envelope. Unknown fields are rejected (strict canonical contract) and
// any trailing content is rejected as 400 (spec §6.4: a single JSON document).
// Input is a json.RawMessage, so any valid JSON value decodes losslessly; schema
// validation is deferred to the execution stage.
//
// The trailing-content check uses a SECOND Decode rather than Decoder.More:
// More only peeks for the next value-start token and treats a stray `]` or `}`
// (or any trailing bytes that are not a value start) as "no more", which would
// let `{envelope} ]` and `{envelope} }` slip through to the stub. A second
// Decode requires the stream to be at exactly EOF afterwards: a nil return
// means another value/document follows, and any non-EOF error means stray
// trailing bytes (`]`, `}`, or garbage), all of which are rejected as 400.
func decodeInvokeEnvelope(data []byte) (invokeRequest, error) {
	var request invokeRequest
	if len(data) == 0 {
		return request, api.NewHTTPErrorWithCode(http.StatusBadRequest, "invalid request body", extensionInvalidBodyCode)
	}
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&request); err != nil {
		return request, api.NewHTTPErrorWithCode(http.StatusBadRequest, "invalid request body", extensionInvalidBodyCode)
	}
	var throwaway json.RawMessage
	if err := decoder.Decode(&throwaway); err != io.EOF {
		return request, api.NewHTTPErrorWithCode(http.StatusBadRequest, "request body must be a single JSON document", extensionTrailingJSONCode)
	}
	if err := validateInvokeEnvelope(request); err != nil {
		return request, err
	}
	return request, nil
}

// validateInvokeEnvelope verifies the canonical envelope's required fields are
// present (spec §6.4): input (the Action payload) must be supplied. A missing
// required field is rejected as 400 rather than reaching the not-wired stub.
// Field-level schema validation of input happens in the execution stage.
func validateInvokeEnvelope(request invokeRequest) error {
	if len(request.Input) == 0 {
		return api.NewHTTPErrorWithCode(http.StatusBadRequest, "input is required", extensionInvalidBodyCode)
	}
	return nil
}

// extensionNotFoundCode is the stable error code for a missing Extension domain
// or action. The error itself is constructed inline at each call site (a shared
// helper would be a one-line thin wrapper).
const extensionNotFoundCode = "not_found"

// extensionBodyTooLargeCode is the stable error code for an invoke body that
// exceeds the transport envelope cap before decode.
const extensionBodyTooLargeCode = "extension_body_too_large"

// extensionInvalidBodyCode is the stable error code for an invoke body that is
// absent or not exactly one valid JSON document.
const extensionInvalidBodyCode = "extension_invalid_body"

// extensionTrailingJSONCode is the stable error code for an invoke body that
// contains trailing JSON after the single envelope document.
const extensionTrailingJSONCode = "extension_trailing_json"