package runner

import (
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"image"
	imagedraw "image/draw"
	"image/jpeg"
	_ "image/png"
	"io"
	"math"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	_ "golang.org/x/image/bmp"
	"golang.org/x/image/draw"
	_ "golang.org/x/image/tiff"
	_ "golang.org/x/image/webp"
)

const (
	maxInputImageConversionBytes  = 32 << 20
	maxInputImageConversionPixels = 100_000_000
)

type inputImageResolutionRange struct {
	MinLong  int
	MinShort int
	MaxLong  int
	MaxShort int
}

type inputImageConstraints struct {
	Resolution inputImageResolutionRange
	MinAspect  float64
	MaxAspect  float64
}

func (s *Service) normalizeVideoInputImages(ctx context.Context, body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
	constraints, ok := candidateInputImageConstraints(candidate)
	if !ok {
		return body, nil
	}
	value, err := s.normalizeVideoInputImageValue(ctx, body, nil, constraints)
	if err != nil {
		return nil, err
	}
	out, _ := value.(map[string]any)
	if out == nil {
		return map[string]any{}, nil
	}
	return out, nil
}

func (s *Service) normalizeVideoInputImageValue(ctx context.Context, value any, path []string, constraints inputImageConstraints) (any, error) {
	switch typed := value.(type) {
	case map[string]any:
		next := make(map[string]any, len(typed))
		for key, item := range typed {
			normalized, err := s.normalizeVideoInputImageValue(ctx, item, append(path, key), constraints)
			if err != nil {
				return nil, err
			}
			next[key] = normalized
		}
		return next, nil
	case []any:
		next := make([]any, 0, len(typed))
		for index, item := range typed {
			normalized, err := s.normalizeVideoInputImageValue(ctx, item, append(path, fmt.Sprintf("[%d]", index)), constraints)
			if err != nil {
				return nil, err
			}
			next = append(next, normalized)
		}
		return next, nil
	case string:
		if !imageInputFieldNeedsHydration(path) {
			return value, nil
		}
		normalized, converted, original, target, err := s.normalizeVideoInputImageSource(ctx, typed, constraints)
		if err != nil {
			param := requestInputImageParam(path)
			return nil, &clients.ClientError{
				Code:       "invalid_parameter",
				Message:    fmt.Sprintf("输入图片 %s 自动转换失败：%s", param, err.Error()),
				Param:      param,
				StatusCode: http.StatusBadRequest,
				Retryable:  false,
			}
		}
		if converted && s.logger != nil {
			s.logger.Info(
				"video input image auto-converted",
				"param", requestInputImageParam(path),
				"original", fmt.Sprintf("%dx%d", original.X, original.Y),
				"converted", fmt.Sprintf("%dx%d", target.X, target.Y),
			)
		}
		return normalized, nil
	default:
		return value, nil
	}
}

func candidateInputImageConstraints(candidate store.RuntimeModelCandidate) (inputImageConstraints, bool) {
	modelTypes := []string{candidate.ModelType, "image_to_video", "omni_video"}
	seen := map[string]bool{}
	for _, modelType := range modelTypes {
		modelType = strings.TrimSpace(modelType)
		if modelType == "" || seen[modelType] {
			continue
		}
		seen[modelType] = true
		values, _ := candidate.Capabilities[modelType].(map[string]any)
		if constraints, ok := parseInputImageConstraints(values); ok {
			return constraints, true
		}
	}
	return inputImageConstraints{}, false
}

func parseInputImageConstraints(values map[string]any) (inputImageConstraints, bool) {
	if values == nil {
		return inputImageConstraints{}, false
	}
	resolution, _ := values["input_image_resolution_range"].(map[string]any)
	minimum, _ := resolution["min"].(map[string]any)
	maximum, _ := resolution["max"].(map[string]any)
	aspect, ok := inputImageNumberPair(values["input_image_aspect_ratio_range"])
	constraints := inputImageConstraints{
		Resolution: inputImageResolutionRange{
			MinLong:  int(math.Round(floatFromAny(minimum["long_edge"]))),
			MinShort: int(math.Round(floatFromAny(minimum["short_edge"]))),
			MaxLong:  int(math.Round(floatFromAny(maximum["long_edge"]))),
			MaxShort: int(math.Round(floatFromAny(maximum["short_edge"]))),
		},
		MinAspect: aspect[0],
		MaxAspect: aspect[1],
	}
	if !ok ||
		constraints.Resolution.MinLong <= 0 ||
		constraints.Resolution.MinShort <= 0 ||
		constraints.Resolution.MaxLong < constraints.Resolution.MinLong ||
		constraints.Resolution.MaxShort < constraints.Resolution.MinShort ||
		constraints.MinAspect <= 0 ||
		constraints.MaxAspect < constraints.MinAspect {
		return inputImageConstraints{}, false
	}
	return constraints, true
}

func inputImageNumberPair(value any) ([2]float64, bool) {
	switch typed := value.(type) {
	case []any:
		if len(typed) != 2 {
			return [2]float64{}, false
		}
		pair := [2]float64{floatFromAny(typed[0]), floatFromAny(typed[1])}
		return pair, pair[0] > 0 && pair[1] > 0
	case []float64:
		if len(typed) != 2 {
			return [2]float64{}, false
		}
		return [2]float64{typed[0], typed[1]}, typed[0] > 0 && typed[1] > 0
	default:
		return [2]float64{}, false
	}
}

func (s *Service) normalizeVideoInputImageSource(ctx context.Context, source string, constraints inputImageConstraints) (string, bool, image.Point, image.Point, error) {
	source = strings.TrimSpace(source)
	if source == "" || strings.HasPrefix(strings.ToLower(source), "asset://") {
		return source, false, image.Point{}, image.Point{}, nil
	}
	payload, err := s.readVideoInputImageBytes(ctx, source)
	if err != nil {
		return "", false, image.Point{}, image.Point{}, err
	}
	config, _, err := image.DecodeConfig(bytes.NewReader(payload))
	if err != nil || config.Width <= 0 || config.Height <= 0 {
		return "", false, image.Point{}, image.Point{}, fmt.Errorf("图片无法解码")
	}
	if config.Width > 50_000 || config.Height > 50_000 || int64(config.Width)*int64(config.Height) > maxInputImageConversionPixels {
		return "", false, image.Point{}, image.Point{}, fmt.Errorf("图片像素数量超过转换上限")
	}
	original := image.Pt(config.Width, config.Height)
	if inputImageWithinConstraints(original, constraints) {
		return source, false, original, original, nil
	}

	release, err := s.acquireImageNormalizationSlot(ctx)
	if err != nil {
		return "", false, original, image.Point{}, err
	}
	defer release()
	decoded, _, err := image.Decode(bytes.NewReader(payload))
	if err != nil {
		return "", false, original, image.Point{}, fmt.Errorf("图片无法解码")
	}
	target := resolveInputImageTarget(original, constraints)
	canvas := image.NewNRGBA(image.Rect(0, 0, target.X, target.Y))
	imagedraw.Draw(canvas, canvas.Bounds(), &image.Uniform{C: averageImageColor(decoded)}, image.Point{}, imagedraw.Src)

	sourceBounds := decoded.Bounds()
	scale := math.Min(float64(target.X)/float64(sourceBounds.Dx()), float64(target.Y)/float64(sourceBounds.Dy()))
	contentWidth := max(1, int(math.Round(float64(sourceBounds.Dx())*scale)))
	contentHeight := max(1, int(math.Round(float64(sourceBounds.Dy())*scale)))
	left := (target.X - contentWidth) / 2
	top := (target.Y - contentHeight) / 2
	draw.ApproxBiLinear.Scale(
		canvas,
		image.Rect(left, top, left+contentWidth, top+contentHeight),
		decoded,
		sourceBounds,
		draw.Over,
		nil,
	)

	var output bytes.Buffer
	if err := jpeg.Encode(&output, canvas, &jpeg.Options{Quality: 90}); err != nil {
		return "", false, original, target, fmt.Errorf("图片编码失败")
	}
	if !inputImageWithinConstraints(target, constraints) {
		return "", false, original, target, fmt.Errorf("转换结果仍不符合平台限制")
	}
	return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(output.Bytes()), true, original, target, nil
}

func (s *Service) acquireImageNormalizationSlot(ctx context.Context) (func(), error) {
	if s.imageNormalizeSlots == nil {
		return func() {}, nil
	}
	select {
	case s.imageNormalizeSlots <- struct{}{}:
		var once sync.Once
		return func() {
			once.Do(func() { <-s.imageNormalizeSlots })
		}, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (s *Service) readVideoInputImageBytes(ctx context.Context, source string) ([]byte, error) {
	lower := strings.ToLower(source)
	if strings.HasPrefix(lower, "data:") || (!strings.Contains(source, "://") && !strings.HasPrefix(source, "/")) {
		payload, err := decodeBase64Payload(source)
		if err != nil {
			return nil, fmt.Errorf("图片数据不是有效 Base64")
		}
		if len(payload) > maxInputImageConversionBytes {
			return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
		}
		return payload, nil
	}
	if localPath := s.localPathFromRequestAssetURL(source); localPath != "" {
		payload, err := os.ReadFile(localPath)
		if err != nil {
			return nil, fmt.Errorf("托管图片读取失败")
		}
		if len(payload) > maxInputImageConversionBytes {
			return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
		}
		return payload, nil
	}
	if !strings.HasPrefix(lower, "http://") && !strings.HasPrefix(lower, "https://") {
		return nil, fmt.Errorf("图片地址格式不受支持")
	}

	requestCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	request, err := http.NewRequestWithContext(requestCtx, http.MethodGet, source, nil)
	if err != nil {
		return nil, fmt.Errorf("图片地址无效")
	}
	httpClient := generatedAssetHTTPClient(false, false)
	httpClient.Timeout = 10 * time.Second
	httpClient.CheckRedirect = func(request *http.Request, via []*http.Request) error {
		if len(via) >= 3 {
			return fmt.Errorf("图片地址重定向次数过多")
		}
		if !requestAssetURLIsPublic("", request.URL.String()) {
			return fmt.Errorf("图片地址重定向到受限网络")
		}
		return nil
	}
	response, err := httpClient.Do(request)
	if err != nil {
		return nil, fmt.Errorf("图片读取失败")
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return nil, fmt.Errorf("图片读取失败，HTTP %d", response.StatusCode)
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, maxInputImageConversionBytes+1))
	if err != nil {
		return nil, fmt.Errorf("图片读取失败")
	}
	if len(payload) > maxInputImageConversionBytes {
		return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
	}
	return payload, nil
}

func inputImageWithinConstraints(size image.Point, constraints inputImageConstraints) bool {
	longEdge := max(size.X, size.Y)
	shortEdge := min(size.X, size.Y)
	ratio := float64(size.X) / float64(size.Y)
	return longEdge >= constraints.Resolution.MinLong &&
		longEdge <= constraints.Resolution.MaxLong &&
		shortEdge >= constraints.Resolution.MinShort &&
		shortEdge <= constraints.Resolution.MaxShort &&
		ratio >= constraints.MinAspect &&
		ratio <= constraints.MaxAspect
}

func resolveInputImageTarget(source image.Point, constraints inputImageConstraints) image.Point {
	sourceRatio := float64(source.X) / float64(source.Y)
	targetRatio := math.Min(constraints.MaxAspect, math.Max(constraints.MinAspect, sourceRatio))
	landscape := targetRatio >= 1
	normalizedRatio := targetRatio
	var sourceShort float64
	if landscape {
		sourceShort = math.Max(float64(source.Y), float64(source.X)/targetRatio)
	} else {
		normalizedRatio = 1 / targetRatio
		sourceShort = math.Max(float64(source.X), float64(source.Y)*targetRatio)
	}

	shortEdge := int(math.Round(sourceShort))
	shortEdge = min(constraints.Resolution.MaxShort, max(constraints.Resolution.MinShort, shortEdge))
	longEdge := int(math.Floor(float64(shortEdge) * normalizedRatio))
	if longEdge > constraints.Resolution.MaxLong {
		longEdge = constraints.Resolution.MaxLong
		shortEdge = int(math.Ceil(float64(longEdge) / normalizedRatio))
	}
	if longEdge < constraints.Resolution.MinLong {
		longEdge = constraints.Resolution.MinLong
		shortEdge = int(math.Ceil(float64(longEdge) / normalizedRatio))
	}
	shortEdge = min(constraints.Resolution.MaxShort, max(constraints.Resolution.MinShort, shortEdge))
	longEdge = int(math.Floor(float64(shortEdge) * normalizedRatio))
	longEdge = min(constraints.Resolution.MaxLong, max(constraints.Resolution.MinLong, longEdge))
	if landscape {
		return image.Pt(longEdge, shortEdge)
	}
	return image.Pt(shortEdge, longEdge)
}

func averageImageColor(source image.Image) imageColor {
	bounds := source.Bounds()
	stepX := max(1, bounds.Dx()/32)
	stepY := max(1, bounds.Dy()/32)
	var red, green, blue, count uint64
	for y := bounds.Min.Y; y < bounds.Max.Y; y += stepY {
		for x := bounds.Min.X; x < bounds.Max.X; x += stepX {
			r, g, b, _ := source.At(x, y).RGBA()
			red += uint64(r >> 8)
			green += uint64(g >> 8)
			blue += uint64(b >> 8)
			count++
		}
	}
	if count == 0 {
		return imageColor{R: 255, G: 255, B: 255, A: 255}
	}
	return imageColor{R: uint8(red / count), G: uint8(green / count), B: uint8(blue / count), A: 255}
}

type imageColor struct {
	R, G, B, A uint8
}

func (c imageColor) RGBA() (r, g, b, a uint32) {
	return uint32(c.R) * 0x101, uint32(c.G) * 0x101, uint32(c.B) * 0x101, uint32(c.A) * 0x101
}

func requestInputImageParam(path []string) string {
	var builder strings.Builder
	for _, segment := range path {
		if strings.HasPrefix(segment, "[") {
			builder.WriteString(segment)
			continue
		}
		if builder.Len() > 0 {
			builder.WriteByte('.')
		}
		builder.WriteString(segment)
	}
	if builder.Len() == 0 {
		return "image"
	}
	return builder.String()
}
