package middleware

import (
	"net/http"
	"strconv"
	"strings"
)

const (
	corsAllowAnyOrigin   = "*"
	corsWildcardMarker   = "*"
	corsPreflightMaxAge  = 300
	corsPreflightSuccess = http.StatusNoContent
)

var (
	corsAllowedMethods = []string{http.MethodGet, http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete, http.MethodOptions}
	corsAllowedHeaders = []string{"Accept", "Authorization", "Content-Type", "X-CSRF-Token"}
	corsExposedHeaders = []string{"Link", "X-Total-Count"}
)

type originWildcard struct {
	prefix string
	suffix string
}

func (w originWildcard) matches(origin string) bool {
	return len(origin) >= len(w.prefix)+len(w.suffix) && strings.HasPrefix(origin, w.prefix) && strings.HasSuffix(origin, w.suffix)
}

type originPolicy struct {
	anyOrigin bool
	exact     map[string]struct{}
	wildcards []originWildcard
}

func newOriginPolicy(origins []string) originPolicy {
	policy := originPolicy{exact: map[string]struct{}{}}
	for _, origin := range origins {
		origin = strings.ToLower(strings.TrimSpace(origin))
		prefix, suffix, hasWildcard := strings.Cut(origin, corsWildcardMarker)
		switch {
		case origin == "":
		case origin == corsAllowAnyOrigin:
			policy.anyOrigin = true
		case hasWildcard:
			policy.wildcards = append(policy.wildcards, originWildcard{prefix: prefix, suffix: suffix})
		default:
			policy.exact[origin] = struct{}{}
		}
	}
	return policy
}

func (p originPolicy) allows(origin string) bool {
	if p.anyOrigin {
		return true
	}
	origin = strings.ToLower(origin)
	if _, ok := p.exact[origin]; ok {
		return true
	}
	for _, wildcard := range p.wildcards {
		if wildcard.matches(origin) {
			return true
		}
	}
	return false
}

func isPreflight(r *http.Request) bool {
	return r.Method == http.MethodOptions && r.Header.Get("Access-Control-Request-Method") != ""
}

func containsFold(allowed []string, value string) bool {
	for _, candidate := range allowed {
		if strings.EqualFold(candidate, value) {
			return true
		}
	}
	return false
}

func requestedHeadersAllowed(r *http.Request) bool {
	requested := r.Header.Get("Access-Control-Request-Headers")
	if requested == "" {
		return true
	}
	for _, header := range strings.Split(requested, ",") {
		if header = strings.TrimSpace(header); header != "" && !containsFold(corsAllowedHeaders, header) {
			return false
		}
	}
	return true
}

func CORS(allowedOrigins []string) func(http.Handler) http.Handler {
	policy := newOriginPolicy(allowedOrigins)

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			origin := r.Header.Get("Origin")
			if origin == "" {
				next.ServeHTTP(w, r)
				return
			}

			w.Header().Add("Vary", "Origin")
			if isPreflight(r) {
				answerPreflight(w, r, origin, policy)
				return
			}

			if policy.allows(origin) {
				w.Header().Set("Access-Control-Allow-Origin", origin)
				w.Header().Set("Access-Control-Expose-Headers", strings.Join(corsExposedHeaders, ", "))
			}
			next.ServeHTTP(w, r)
		})
	}
}

func answerPreflight(w http.ResponseWriter, r *http.Request, origin string, policy originPolicy) {
	w.Header().Add("Vary", "Access-Control-Request-Method")
	w.Header().Add("Vary", "Access-Control-Request-Headers")

	method := r.Header.Get("Access-Control-Request-Method")
	if !policy.allows(origin) || !containsFold(corsAllowedMethods, method) || !requestedHeadersAllowed(r) {
		w.WriteHeader(http.StatusOK)
		return
	}

	w.Header().Set("Access-Control-Allow-Origin", origin)
	w.Header().Set("Access-Control-Allow-Methods", strings.ToUpper(method))
	w.Header().Set("Access-Control-Allow-Headers", strings.Join(corsAllowedHeaders, ", "))
	w.Header().Set("Access-Control-Max-Age", strconv.Itoa(corsPreflightMaxAge))
	w.WriteHeader(corsPreflightSuccess)
}