// Изменено: 2026-04-09
// Auth middleware для shared-sqs: извлекает AccessKeyId из AWS Authorization header
// и помещает найденного тенанта в context запроса.
package auth

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"shared-sqs/app/tenant"
)

// TenantContextKey — ключ для хранения тенанта в request context.
// Тип contextKey предотвращает конфликты с другими пакетами.
type contextKey string

const TenantContextKey contextKey = "tenant"

const (
	sigV4DateFormat = "20060102T150405Z"
	maxClockSkew    = 15 * time.Minute
)

// AuthMiddleware — middleware: ищет тенанта по AccessKeyId из AWS Authorization header.
// Пропускает /health и /admin/** без tenant-аутентификации.
// Ловушка #4: не ставим короткий таймаут — ReceiveMessage с long polling держит соединение до 20 сек.
func AuthMiddleware(store *tenant.TenantStore) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// /health — без auth
			if r.URL.Path == "/health" {
				next.ServeHTTP(w, r)
				return
			}
			// /admin/** — отдельная auth (bearer token, см. admin_handlers.go)
			if strings.HasPrefix(r.URL.Path, "/admin/") {
				next.ServeHTTP(w, r)
				return
			}

			accessKeyID := extractAccessKeyID(r)
			if accessKeyID == "" {
				writeSQSAuthError(w, "MissingAuthenticationToken", "Request must contain either AccessKeyId or X-Amz-Credential")
				return
			}

			t, ok := store.GetByAccessKey(accessKeyID)
			if !ok || !t.Active {
				writeSQSAuthError(w, "InvalidClientTokenId", "The security token included in the request is invalid")
				return
			}

			if err := verifySigV4(r, t); err != nil {
				writeSQSAuthError(w, "SignatureDoesNotMatch", err.Error())
				return
			}

			ctx := context.WithValue(r.Context(), TenantContextKey, t)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// extractAccessKeyID — извлекает AWS AccessKeyId из запроса.
// Поддерживает оба варианта: Authorization header (Signature V4) и X-Amz-Credential query param (presigned URLs).
// Ловушка #3: AWS CLI ВСЕГДА отправляет Signature V4 — нужно парсить, даже не проверяя подпись.
// Ловушка #5: X-Amz-Security-Token (STS) — игнорируем.
func extractAccessKeyID(r *http.Request) string {
	// Вариант 1: Authorization header
	// Формат: "AWS4-HMAC-SHA256 Credential={AccessKeyId}/{date}/{region}/sqs/aws4_request, ..."
	auth := r.Header.Get("Authorization")
	if strings.HasPrefix(auth, "AWS4-HMAC-SHA256") {
		idx := strings.Index(auth, "Credential=")
		if idx >= 0 {
			rest := auth[idx+len("Credential="):]
			slashIdx := strings.Index(rest, "/")
			if slashIdx > 0 {
				return rest[:slashIdx]
			}
		}
	}

	// Вариант 2: Query parameter (presigned URLs)
	// Формат: X-Amz-Credential={AccessKeyId}/{date}/{region}/sqs/aws4_request
	if cred := r.URL.Query().Get("X-Amz-Credential"); cred != "" {
		parts := strings.SplitN(cred, "/", 2)
		if len(parts) > 0 && parts[0] != "" {
			return parts[0]
		}
	}

	return ""
}

// verifySigV4 — полноценная проверка подписи AWS SigV4 для header и presigned запросов.
func verifySigV4(r *http.Request, t *tenant.Tenant) error {
	authHeader := r.Header.Get("Authorization")
	if strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256") {
		return verifyHeaderSigV4(r, t, authHeader)
	}
	if r.URL.Query().Get("X-Amz-Credential") != "" {
		return verifyPresignedSigV4(r, t)
	}
	return errors.New("missing SigV4 authorization")
}

func verifyHeaderSigV4(r *http.Request, t *tenant.Tenant, authHeader string) error {
	credential, signedHeaders, providedSignature, err := parseAuthorizationHeader(authHeader)
	if err != nil {
		return err
	}
	accessKey, shortDate, region, service, terminal, err := parseCredentialScope(credential)
	if err != nil {
		return err
	}
	if accessKey != t.AccessKey {
		return errors.New("access key mismatch")
	}
	if service != "sqs" || terminal != "aws4_request" {
		return errors.New("invalid credential scope")
	}

	amzDate := r.Header.Get("X-Amz-Date")
	if amzDate == "" {
		return errors.New("missing X-Amz-Date header")
	}
	requestTime, err := time.Parse(sigV4DateFormat, amzDate)
	if err != nil {
		return fmt.Errorf("invalid X-Amz-Date: %w", err)
	}
	if absDuration(time.Now().UTC().Sub(requestTime)) > maxClockSkew {
		return errors.New("request time skew is too large")
	}

	body, err := readAndRestoreBody(r)
	if err != nil {
		return err
	}
	payloadHash := r.Header.Get("X-Amz-Content-Sha256")
	if payloadHash == "" {
		payloadHash = sha256Hex(body)
	}

	canonicalRequest, err := buildCanonicalRequest(r, signedHeaders, payloadHash, true)
	if err != nil {
		return err
	}
	credentialScope := strings.Join([]string{shortDate, region, service, terminal}, "/")
	stringToSign := buildStringToSign(amzDate, credentialScope, canonicalRequest)
	expectedSignature := calculateSigV4Signature(t.SecretKey, shortDate, region, service, stringToSign)

	if subtle.ConstantTimeCompare([]byte(strings.ToLower(expectedSignature)), []byte(strings.ToLower(providedSignature))) != 1 {
		return errors.New("signature mismatch")
	}
	return nil
}

func verifyPresignedSigV4(r *http.Request, t *tenant.Tenant) error {
	query := r.URL.Query()
	if query.Get("X-Amz-Algorithm") != "AWS4-HMAC-SHA256" {
		return errors.New("unsupported X-Amz-Algorithm")
	}
	credential := query.Get("X-Amz-Credential")
	if credential == "" {
		return errors.New("missing X-Amz-Credential")
	}
	providedSignature := query.Get("X-Amz-Signature")
	if providedSignature == "" {
		return errors.New("missing X-Amz-Signature")
	}
	signedHeaders := query.Get("X-Amz-SignedHeaders")
	if signedHeaders == "" {
		return errors.New("missing X-Amz-SignedHeaders")
	}
	amzDate := query.Get("X-Amz-Date")
	if amzDate == "" {
		return errors.New("missing X-Amz-Date")
	}
	expiresSeconds, err := strconv.Atoi(query.Get("X-Amz-Expires"))
	if err != nil || expiresSeconds < 1 || expiresSeconds > 604800 {
		return errors.New("invalid X-Amz-Expires")
	}

	requestTime, err := time.Parse(sigV4DateFormat, amzDate)
	if err != nil {
		return fmt.Errorf("invalid X-Amz-Date: %w", err)
	}
	if time.Now().UTC().After(requestTime.Add(time.Duration(expiresSeconds) * time.Second)) {
		return errors.New("presigned request has expired")
	}

	accessKey, shortDate, region, service, terminal, err := parseCredentialScope(credential)
	if err != nil {
		return err
	}
	if accessKey != t.AccessKey {
		return errors.New("access key mismatch")
	}
	if service != "sqs" || terminal != "aws4_request" {
		return errors.New("invalid credential scope")
	}

	body, err := readAndRestoreBody(r)
	if err != nil {
		return err
	}
	payloadHash := query.Get("X-Amz-Content-Sha256")
	if payloadHash == "" {
		payloadHash = r.Header.Get("X-Amz-Content-Sha256")
	}
	if payloadHash == "" {
		payloadHash = "UNSIGNED-PAYLOAD"
		if len(body) > 0 {
			payloadHash = sha256Hex(body)
		}
	}

	canonicalRequest, err := buildCanonicalRequest(r, signedHeaders, payloadHash, false)
	if err != nil {
		return err
	}
	credentialScope := strings.Join([]string{shortDate, region, service, terminal}, "/")
	stringToSign := buildStringToSign(amzDate, credentialScope, canonicalRequest)
	expectedSignature := calculateSigV4Signature(t.SecretKey, shortDate, region, service, stringToSign)

	if subtle.ConstantTimeCompare([]byte(strings.ToLower(expectedSignature)), []byte(strings.ToLower(providedSignature))) != 1 {
		return errors.New("signature mismatch")
	}
	return nil
}

func parseAuthorizationHeader(authHeader string) (credential string, signedHeaders string, signature string, err error) {
	if !strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256") {
		return "", "", "", errors.New("unsupported authorization algorithm")
	}
	fields := strings.Split(strings.TrimSpace(strings.TrimPrefix(authHeader, "AWS4-HMAC-SHA256")), ",")
	parts := map[string]string{}
	for _, field := range fields {
		kv := strings.SplitN(strings.TrimSpace(field), "=", 2)
		if len(kv) != 2 {
			continue
		}
		parts[kv[0]] = kv[1]
	}
	credential = parts["Credential"]
	signedHeaders = strings.ToLower(parts["SignedHeaders"])
	signature = parts["Signature"]
	if credential == "" || signedHeaders == "" || signature == "" {
		return "", "", "", errors.New("malformed Authorization header")
	}
	return credential, signedHeaders, signature, nil
}

func parseCredentialScope(credential string) (accessKey, shortDate, region, service, terminal string, err error) {
	parts := strings.Split(credential, "/")
	if len(parts) != 5 {
		return "", "", "", "", "", errors.New("invalid Credential scope")
	}
	return parts[0], parts[1], parts[2], parts[3], parts[4], nil
}

func buildCanonicalRequest(r *http.Request, signedHeaders, payloadHash string, includeSignatureQuery bool) (string, error) {
	canonicalURI := r.URL.EscapedPath()
	if canonicalURI == "" {
		canonicalURI = "/"
	}
	canonicalQuery := canonicalizeQueryString(r.URL.Query(), includeSignatureQuery)

	headers := strings.Split(strings.ToLower(signedHeaders), ";")
	canonicalHeaders := strings.Builder{}
	cleanSignedHeaders := make([]string, 0, len(headers))
	for _, h := range headers {
		name := strings.TrimSpace(h)
		if name == "" {
			continue
		}
		value, ok := canonicalHeaderValue(r, name)
		if !ok {
			return "", fmt.Errorf("missing signed header: %s", name)
		}
		canonicalHeaders.WriteString(name)
		canonicalHeaders.WriteString(":")
		canonicalHeaders.WriteString(normalizeHeaderSpace(value))
		canonicalHeaders.WriteString("\n")
		cleanSignedHeaders = append(cleanSignedHeaders, name)
	}
	if len(cleanSignedHeaders) == 0 {
		return "", errors.New("empty signed headers")
	}

	return strings.Join([]string{
		r.Method,
		canonicalURI,
		canonicalQuery,
		canonicalHeaders.String(),
		strings.Join(cleanSignedHeaders, ";"),
		payloadHash,
	}, "\n"), nil
}

func canonicalHeaderValue(r *http.Request, name string) (string, bool) {
	if name == "host" {
		h := r.Host
		if h == "" {
			h = r.URL.Host
		}
		return h, h != ""
	}
	key := http.CanonicalHeaderKey(name)
	values, ok := r.Header[key]
	if !ok || len(values) == 0 {
		return "", false
	}
	return strings.Join(values, ","), true
}

func canonicalizeQueryString(query url.Values, includeSignature bool) string {
	pairs := make([]string, 0)
	for key, values := range query {
		if !includeSignature && strings.EqualFold(key, "X-Amz-Signature") {
			continue
		}
		sortedValues := append([]string(nil), values...)
		sort.Strings(sortedValues)
		if len(sortedValues) == 0 {
			pairs = append(pairs, awsPercentEncode(key)+"=")
			continue
		}
		for _, value := range sortedValues {
			pairs = append(pairs, awsPercentEncode(key)+"="+awsPercentEncode(value))
		}
	}
	sort.Strings(pairs)
	return strings.Join(pairs, "&")
}

func awsPercentEncode(s string) string {
	encoded := url.QueryEscape(s)
	encoded = strings.ReplaceAll(encoded, "+", "%20")
	encoded = strings.ReplaceAll(encoded, "*", "%2A")
	encoded = strings.ReplaceAll(encoded, "%7E", "~")
	return encoded
}

func normalizeHeaderSpace(v string) string {
	return strings.Join(strings.Fields(strings.TrimSpace(v)), " ")
}

func buildStringToSign(amzDate, credentialScope, canonicalRequest string) string {
	canonicalHash := sha256.Sum256([]byte(canonicalRequest))
	return strings.Join([]string{
		"AWS4-HMAC-SHA256",
		amzDate,
		credentialScope,
		hex.EncodeToString(canonicalHash[:]),
	}, "\n")
}

func calculateSigV4Signature(secretKey, shortDate, region, service, stringToSign string) string {
	kDate := hmacSHA256([]byte("AWS4"+secretKey), shortDate)
	kRegion := hmacSHA256(kDate, region)
	kService := hmacSHA256(kRegion, service)
	kSigning := hmacSHA256(kService, "aws4_request")
	sig := hmacSHA256(kSigning, stringToSign)
	return hex.EncodeToString(sig)
}

func hmacSHA256(key []byte, data string) []byte {
	m := hmac.New(sha256.New, key)
	_, _ = m.Write([]byte(data))
	return m.Sum(nil)
}

func readAndRestoreBody(r *http.Request) ([]byte, error) {
	if r.Body == nil {
		return nil, nil
	}
	body, err := io.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("read request body: %w", err)
	}
	r.Body = io.NopCloser(bytes.NewReader(body))
	return body, nil
}

func sha256Hex(data []byte) string {
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:])
}

func absDuration(d time.Duration) time.Duration {
	if d < 0 {
		return -d
	}
	return d
}

// sqsAuthError — AWS-совместимый XML ответ об ошибке аутентификации.
type sqsAuthError struct {
	XMLName   xml.Name     `xml:"ErrorResponse"`
	Error     sqsErrorBody `xml:"Error"`
	RequestID string       `xml:"RequestId"`
}

type sqsErrorBody struct {
	Type    string `xml:"Type"`
	Code    string `xml:"Code"`
	Message string `xml:"Message"`
}

// writeSQSAuthError — отвечает AWS-совместимым XML с кодом 403.
func writeSQSAuthError(w http.ResponseWriter, code, message string) {
	w.Header().Set("Content-Type", "application/xml")
	w.WriteHeader(http.StatusForbidden)
	resp := sqsAuthError{
		Error: sqsErrorBody{
			Type:    "Sender",
			Code:    code,
			Message: message,
		},
		RequestID: "00000000-0000-0000-0000-000000000000",
	}
	data, _ := xml.Marshal(resp)
	w.Write(data)
}