package egress

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime"
	"net/http"
	"strings"
)

const (
	MaxRawResponseBytes   = 1 << 20
	MaxFinalResponseBytes = 64 << 10
)

var (
	ErrRawResponseTooLarge    = errors.New("extension upstream response exceeds 1 MiB")
	ErrFinalResponseTooLarge  = errors.New("extension final response exceeds 64 KiB")
	ErrUnsupportedContentType = errors.New(
		"extension upstream content type is not allowed",
	)
	ErrUnexpectedStatus = errors.New("extension upstream returned an unexpected status")
)

// UpstreamStatusError preserves only the numeric provider status needed for
// stable classification (401 refresh, 429 rate limit, other upstream failure).
// It never carries headers, response body, or the credential-bearing request
// URL.
type UpstreamStatusError struct {
	Status int
}

func (err *UpstreamStatusError) Error() string {
	return fmt.Sprintf("%s: %s", ErrUnexpectedStatus.Error(), statusClass(err.Status))
}

func (err *UpstreamStatusError) Unwrap() error {
	return ErrUnexpectedStatus
}

// UpstreamStatus extracts a sanitized upstream HTTP status from err.
func UpstreamStatus(err error) (int, bool) {
	var statusErr *UpstreamStatusError
	if !errors.As(err, &statusErr) || statusErr == nil {
		return 0, false
	}
	return statusErr.Status, true
}

// Result is the bounded, normalized provider result. Data is exactly one JSON
// value; text responses are represented as a JSON string.
type Result struct {
	Data           json.RawMessage
	UpstreamStatus int
	StatusClass    string
	ContentType    string
	RequestBytes   int64
	ResponseBytes  int64
	FinalBytes     int
}

func normalizeResponse(
	response *http.Response,
	requestBytes int64,
	allowedStatuses []int,
	allowedContentTypes []string,
) (Result, error) {
	if response == nil || response.Body == nil {
		return Result{}, errors.New("extension upstream response is unavailable")
	}
	raw, err := io.ReadAll(io.LimitReader(response.Body, MaxRawResponseBytes+1))
	if err != nil {
		return Result{}, fmt.Errorf("read extension upstream response: %w", err)
	}
	result := Result{
		UpstreamStatus: response.StatusCode,
		StatusClass:    statusClass(response.StatusCode),
		RequestBytes:   requestBytes,
		ResponseBytes:  int64(len(raw)),
	}
	if len(raw) > MaxRawResponseBytes {
		return result, ErrRawResponseTooLarge
	}
	if !containsStatus(allowedStatuses, response.StatusCode) {
		return result, &UpstreamStatusError{Status: response.StatusCode}
	}
	contentType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type"))
	if err != nil {
		return result, ErrUnsupportedContentType
	}
	contentType = strings.ToLower(contentType)
	if !containsContentType(allowedContentTypes, contentType) || !globallyAllowedContentType(contentType) {
		return result, ErrUnsupportedContentType
	}
	result.ContentType = contentType

	var data json.RawMessage
	if contentType == "application/json" ||
		(strings.HasPrefix(contentType, "application/") && strings.HasSuffix(contentType, "+json")) {
		value, err := decodeOneJSON(raw)
		if err != nil {
			return result, fmt.Errorf("decode extension upstream JSON: %w", err)
		}
		data, err = json.Marshal(value)
		if err != nil {
			return result, fmt.Errorf("normalize extension upstream JSON: %w", err)
		}
	} else {
		encoded, err := json.Marshal(string(bytes.ToValidUTF8(raw, []byte("\uFFFD"))))
		if err != nil {
			return result, fmt.Errorf("normalize extension upstream text: %w", err)
		}
		data = encoded
	}
	probe, err := json.Marshal(struct {
		Data           json.RawMessage `json:"data"`
		UpstreamStatus int             `json:"upstreamStatus"`
	}{
		Data: data, UpstreamStatus: response.StatusCode,
	})
	if err != nil {
		return result, fmt.Errorf("measure extension final response: %w", err)
	}
	if len(probe) > MaxFinalResponseBytes {
		return result, ErrFinalResponseTooLarge
	}
	result.Data = data
	result.FinalBytes = len(probe)
	return result, nil
}

func decodeOneJSON(raw []byte) (any, error) {
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	var value any
	if err := decoder.Decode(&value); err != nil {
		return nil, err
	}
	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		return nil, errors.New("response must be exactly one JSON document")
	}
	return value, nil
}

func containsStatus(allowed []int, status int) bool {
	for _, candidate := range allowed {
		if candidate == status {
			return true
		}
	}
	return false
}

func containsContentType(allowed []string, contentType string) bool {
	for _, candidate := range allowed {
		mediaType, _, err := mime.ParseMediaType(candidate)
		if err == nil && strings.EqualFold(mediaType, contentType) {
			return true
		}
	}
	return false
}

func globallyAllowedContentType(contentType string) bool {
	switch {
	case contentType == "application/json":
		return true
	case strings.HasPrefix(contentType, "application/") && strings.HasSuffix(contentType, "+json"):
		return true
	case strings.HasPrefix(contentType, "text/") && contentType != "text/event-stream":
		return true
	default:
		return false
	}
}

func statusClass(status int) string {
	if status < 100 || status > 999 {
		return "unknown"
	}
	return fmt.Sprintf("%dxx", status/100)
}
