package server

import (
	"fmt"
	"mime"
	"net/http"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
)

var (
	// NOTE: these are just the default allowed response headers.
	// Users are able to extend it with extra allowed response headers via KAS configuration.
	defaultAllowedResponseHeaders = sets.New(
		httpz.DateHeader,
		httpz.VaryHeader,
		httpz.ViaHeader,
		httpz.EtagHeader,
		httpz.ContentTypeHeader,
		httpz.ContentLengthHeader,
		httpz.ContentEncodingHeader,
		httpz.AcceptRangesHeader,
		httpz.ContentRangeHeader,
		httpz.LastModifiedHeader,
		httpz.RetryAfterHeader,
		httpz.LocationHeader,
		httpz.ConnectionHeader,
		httpz.UpgradeHeader,
		httpz.WarningHeader, // deprecated in HTTP standard, but used in Kubernetes

		// WebSocket
		httpz.SecWebSocketProtocolHeader,
		httpz.SecWebSocketExtensionsHeader,
		httpz.SecWebSocketAcceptHeader,
		httpz.SecWebSocketVersionHeader,

		// Kubernetes specific headers
		httpz.AuditIDHeader,
		httpz.XStreamProtocolVersionHeader,
		httpz.XAcceptedStreamProtocolVersionsHeader,

		// GitLab Agent (agentk) headers
		httpz.GitlabAgentVersionHeader,
	)

	allowedResponseHeaderNamePrefixes = []string{
		"X-Kubernetes-", // e.g. Flow Control headers, like X-Kubernetes-Pf-Flowschema-Uid and X-Kubernetes-Pf-Prioritylevel-Uid
	}

	allowedResponseContentTypes = []string{
		runtime.ContentTypeJSON,
		runtime.ContentTypeYAML,
		runtime.ContentTypeProtobuf,
		runtime.ContentTypeCBOR,
		"text/plain",
	}

	// knownBlockedHeaders contains a set of headers that are known to be blocked, but usually sent in regular Kubernetes workflows.
	// This is mostly to reduce noise when reporting on blocked headers.
	knownBlockedHeaders = sets.New(
		httpz.CacheControlHeader,
		httpz.ExpiresHeader,
		httpz.StrictTransportSecurityHeader,
		httpz.XContentTypeOptionsHeader,
		httpz.PragmaHeader,
	)
)

func checkContentType(h http.Header, allowed ...string) error {
	// There should be at most one Content-Type header, but it's not our job to do something about it if there is more.
	// We just ensure thy are all allowed.
nextContentType:
	for _, ct := range h[httpz.ContentTypeHeader] {
		mediatype, _, err := mime.ParseMediaType(ct)
		if err != nil && err != mime.ErrInvalidMediaParameter {
			// Parsing error and not a MIME parameter parsing error, which we ignore.
			return fmt.Errorf("check Content-Type: %w", err)
		}
		for _, a := range allowed {
			if mediatype == a {
				// This one is allowed, onto the next Content-Type header
				continue nextContentType
			}
		}
		return fmt.Errorf("%s not allowed: %s", httpz.ContentTypeHeader, mediatype)
	}
	return nil
}
