package clients

import (
	"math"
	"strconv"
	"strings"
)

const (
	openAIGPTImage2MaxEdge   = 3840
	openAIGPTImage2MinPixels = 655360
	openAIGPTImage2MaxPixels = 8294400
	openAIGPTImage2MaxRatio  = 3
	openAIGPTImage2Multiple  = 16
)

type providerImageDimensions struct {
	width  int
	height int
}

func normalizedProviderAspectRatio(value any) string {
	compact := strings.Map(func(r rune) rune {
		switch r {
		case ' ', '\t', '\r', '\n':
			return -1
		default:
			return r
		}
	}, strings.TrimSpace(stringFromAny(value)))
	parts := strings.Split(compact, ":")
	if len(parts) != 2 {
		return ""
	}
	width, widthErr := strconv.ParseFloat(parts[0], 64)
	height, heightErr := strconv.ParseFloat(parts[1], 64)
	if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
		return ""
	}
	return parts[0] + ":" + parts[1]
}

func normalizedProviderImageResolution(value any) string {
	text := strings.ToLower(strings.TrimSpace(stringFromAny(value)))
	if len(text) < 2 || text[len(text)-1] != 'k' {
		return ""
	}
	magnitude, err := strconv.Atoi(text[:len(text)-1])
	if err != nil || magnitude <= 0 {
		return ""
	}
	return strconv.Itoa(magnitude) + "K"
}

func normalizeOfficialOpenAIImageSize(body map[string]any, originalBody map[string]any) {
	if originalBody == nil {
		originalBody = body
	}
	if width, height, ok := providerPixelDimensions(stringFromAny(originalBody["size"])); ok {
		body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
		return
	}

	model := strings.ToLower(strings.TrimSpace(stringFromAny(body["model"])))
	resolution := normalizedProviderImageResolution(firstNonEmptyString(
		originalBody["resolution"],
		originalBody["size"],
		body["resolution"],
		body["size"],
	))
	if resolution != "" {
		width, height, ok := providerBodyPixelDimensions(body)
		if !ok {
			width, height, ok = providerResolutionDimensions(
				resolution,
				firstNonEmptyString(originalBody["aspect_ratio"], originalBody["aspectRatio"], originalBody["ratio"]),
			)
		}
		if !ok {
			delete(body, "size")
			return
		}
		switch {
		case strings.HasPrefix(model, "gpt-image-2"):
			dimensions := normalizeOpenAIGPTImage2Dimensions(width, height)
			body["size"] = strconv.Itoa(dimensions.width) + "x" + strconv.Itoa(dimensions.height)
		case strings.HasPrefix(model, "gpt-image-"):
			body["size"] = officialOpenAILegacyImageSize(width, height)
		default:
			body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
		}
		return
	}

	if !strings.HasPrefix(model, "gpt-image-") {
		return
	}
	size := strings.ToLower(strings.TrimSpace(stringFromAny(body["size"])))
	if size == "auto" {
		body["size"] = "auto"
		return
	}
	if strings.HasPrefix(model, "gpt-image-2") {
		if width, height, ok := providerPixelDimensions(size); ok {
			body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
		}
		return
	}

	switch size {
	case "1024x1024", "1536x1024", "1024x1536":
		body["size"] = size
		return
	}
	width, height, ok := requestedProviderDimensions(body)
	if !ok {
		body["size"] = "auto"
		return
	}
	body["size"] = officialOpenAILegacyImageSize(width, height)
}

func officialOpenAILegacyImageSize(width int, height int) string {
	switch {
	case width > height:
		return "1536x1024"
	case height > width:
		return "1024x1536"
	default:
		return "1024x1024"
	}
}

func providerBodyPixelDimensions(body map[string]any) (int, int, bool) {
	if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
		return width, height, true
	}
	width := intFromAny(body["width"])
	height := intFromAny(body["height"])
	return width, height, width > 0 && height > 0
}

func providerResolutionDimensions(resolution string, aspectRatio any) (int, int, bool) {
	resolution = normalizedProviderImageResolution(resolution)
	if resolution == "" {
		return 0, 0, false
	}
	magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
	if err != nil || magnitude <= 0 {
		return 0, 0, false
	}
	baseSide := magnitude * 1024
	ratio := providerAspectRatioNumber(aspectRatio)
	width, height := baseSide, baseSide
	if ratio > 1 {
		height = max(1, int(math.Round(float64(baseSide)/ratio)))
	} else if ratio > 0 && ratio < 1 {
		width = max(1, int(math.Round(float64(baseSide)*ratio)))
	}
	return width, height, true
}

func requestedProviderDimensions(body map[string]any) (int, int, bool) {
	if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
		return width, height, true
	}
	width := intFromAny(body["width"])
	height := intFromAny(body["height"])
	if width > 0 && height > 0 {
		return width, height, true
	}
	ratio := providerAspectRatioNumber(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
	if ratio > 1 {
		return 2, 1, true
	}
	if ratio > 0 && ratio < 1 {
		return 1, 2, true
	}
	if ratio == 1 {
		return 1, 1, true
	}
	return 0, 0, false
}

func providerPixelDimensions(value string) (int, int, bool) {
	parts := strings.Split(strings.ToLower(strings.TrimSpace(value)), "x")
	if len(parts) != 2 {
		return 0, 0, false
	}
	width, widthErr := strconv.Atoi(strings.TrimSpace(parts[0]))
	height, heightErr := strconv.Atoi(strings.TrimSpace(parts[1]))
	if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
		return 0, 0, false
	}
	return width, height, true
}

func providerAspectRatioNumber(value any) float64 {
	normalized := normalizedProviderAspectRatio(value)
	if normalized == "" {
		return 0
	}
	parts := strings.Split(normalized, ":")
	width, _ := strconv.ParseFloat(parts[0], 64)
	height, _ := strconv.ParseFloat(parts[1], 64)
	return width / height
}

func normalizeOpenAIGPTImage2Dimensions(width int, height int) providerImageDimensions {
	if width <= 0 || height <= 0 {
		return providerImageDimensions{}
	}
	rawRatio := float64(width) / float64(height)
	ratio := rawRatio
	if ratio < 1 {
		ratio = 1 / ratio
	}
	ratio = min(ratio, float64(openAIGPTImage2MaxRatio))
	landscape := rawRatio >= 1
	longEdge := min(max(width, height), openAIGPTImage2MaxEdge)
	dimensions := buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
	for dimensions.width*dimensions.height > openAIGPTImage2MaxPixels {
		longEdge = max(openAIGPTImage2Multiple, max(dimensions.width, dimensions.height)-openAIGPTImage2Multiple)
		dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
	}
	for dimensions.width*dimensions.height < openAIGPTImage2MinPixels && max(dimensions.width, dimensions.height) < openAIGPTImage2MaxEdge {
		longEdge = min(openAIGPTImage2MaxEdge, max(dimensions.width, dimensions.height)+openAIGPTImage2Multiple)
		dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, true)
	}
	return dimensions
}

func buildOpenAIGPTImage2Dimensions(longEdge int, ratio float64, landscape bool, roundUp bool) providerImageDimensions {
	normalizedLongEdge := roundOpenAIGPTImage2Multiple(float64(longEdge), roundUp)
	normalizedShortEdge := roundOpenAIGPTImage2Multiple(float64(normalizedLongEdge)/ratio, roundUp)
	if float64(normalizedLongEdge)/float64(normalizedShortEdge) > openAIGPTImage2MaxRatio {
		normalizedShortEdge += openAIGPTImage2Multiple
	}
	if landscape {
		return providerImageDimensions{width: normalizedLongEdge, height: normalizedShortEdge}
	}
	return providerImageDimensions{width: normalizedShortEdge, height: normalizedLongEdge}
}

func roundOpenAIGPTImage2Multiple(value float64, roundUp bool) int {
	rounded := math.Round(value / openAIGPTImage2Multiple)
	if roundUp {
		rounded = math.Ceil(value / openAIGPTImage2Multiple)
	}
	return max(openAIGPTImage2Multiple, int(rounded)*openAIGPTImage2Multiple)
}