// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sdk

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"mime"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const maxImageInputBytes = 20 * 1024 * 1024

var imageHTTPClient = &http.Client{
	Timeout: 15 * time.Second,
	Transport: &http.Transport{
		DialContext: dialPublicImageHost,
	},
	CheckRedirect: func(req *http.Request, via []*http.Request) error {
		if len(via) >= 5 {
			return fmt.Errorf("too many redirects")
		}
		if req.URL.Scheme != "https" {
			return fmt.Errorf("image redirect must use https")
		}
		return validateRemoteImageHost(req.Context(), req.URL.Hostname())
	},
}

func normalizeImageURL(ctx context.Context, rawURL string) (string, error) {
	if strings.HasPrefix(rawURL, "data:image/") {
		if err := validateDataImageURL(rawURL); err != nil {
			return "", err
		}
		return rawURL, nil
	}
	return downloadImageAsDataURL(ctx, rawURL)
}

func validateDataImageURL(dataURL string) error {
	header, encoded, ok := strings.Cut(dataURL, ",")
	if !ok {
		return fmt.Errorf("image data URL must contain base64 data")
	}
	if !strings.HasPrefix(header, "data:image/") || !strings.Contains(header, ";base64") {
		return fmt.Errorf("image URL must be a base64 data URL")
	}
	mediaType := strings.TrimPrefix(strings.Split(header, ";")[0], "data:")
	if !isAllowedImageMediaType(mediaType) {
		return fmt.Errorf("unsupported image media type %q", mediaType)
	}
	decoded, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return fmt.Errorf("invalid image base64 data: %w", err)
	}
	if len(decoded) > maxImageInputBytes {
		return fmt.Errorf("image exceeds maximum size of %d bytes", maxImageInputBytes)
	}
	return nil
}

func downloadImageAsDataURL(ctx context.Context, rawURL string) (string, error) {
	parsed, err := url.Parse(rawURL)
	if err != nil {
		return "", fmt.Errorf("invalid image URL: %w", err)
	}
	if parsed.Scheme != "https" {
		return "", fmt.Errorf("image URL must use https")
	}
	if err := validateRemoteImageHost(ctx, parsed.Hostname()); err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create image request: %w", err)
	}
	resp, err := imageHTTPClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to download image: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", fmt.Errorf("image URL returned status %d", resp.StatusCode)
	}
	if resp.ContentLength > maxImageInputBytes {
		return "", fmt.Errorf("image exceeds maximum size of %d bytes", maxImageInputBytes)
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, maxImageInputBytes+1))
	if err != nil {
		return "", fmt.Errorf("failed to read image response: %w", err)
	}
	if len(body) > maxImageInputBytes {
		return "", fmt.Errorf("image exceeds maximum size of %d bytes", maxImageInputBytes)
	}
	mediaType := imageMediaType(resp.Header.Get("Content-Type"), body)
	if !isAllowedImageMediaType(mediaType) {
		return "", fmt.Errorf("unsupported image media type %q", mediaType)
	}
	return fmt.Sprintf("data:%s;base64,%s", mediaType, base64.StdEncoding.EncodeToString(body)), nil
}

func validateRemoteImageHost(ctx context.Context, host string) error {
	if host == "" {
		return fmt.Errorf("image URL host is required")
	}
	_, err := resolvePublicImageIP(ctx, host)
	return err
}

func dialPublicImageHost(ctx context.Context, network, address string) (net.Conn, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	ip, err := resolvePublicImageIP(ctx, host)
	if err != nil {
		return nil, err
	}
	var dialer net.Dialer
	return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
}

func resolvePublicImageIP(ctx context.Context, host string) (net.IP, error) {
	if ip := net.ParseIP(host); ip != nil {
		if isBlockedImageIP(ip) {
			return nil, fmt.Errorf("image URL host resolves to disallowed address")
		}
		return ip, nil
	}
	addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve image URL host: %w", err)
	}
	if len(addrs) == 0 {
		return nil, fmt.Errorf("image URL host has no addresses")
	}
	var firstPublic net.IP
	for _, addr := range addrs {
		if isBlockedImageIP(addr.IP) {
			return nil, fmt.Errorf("image URL host resolves to disallowed address")
		}
		if firstPublic == nil {
			firstPublic = addr.IP
		}
	}
	return firstPublic, nil
}

func isBlockedImageIP(ip net.IP) bool {
	return ip.IsLoopback() ||
		ip.IsPrivate() ||
		ip.IsLinkLocalUnicast() ||
		ip.IsLinkLocalMulticast() ||
		ip.IsMulticast() ||
		ip.IsUnspecified()
}

func imageMediaType(contentType string, body []byte) string {
	mediaType, _, err := mime.ParseMediaType(contentType)
	if err == nil && mediaType != "" {
		return strings.ToLower(mediaType)
	}
	return strings.ToLower(http.DetectContentType(body))
}

func isAllowedImageMediaType(mediaType string) bool {
	switch strings.ToLower(mediaType) {
	case "image/png", "image/jpeg", "image/webp", "image/gif":
		return true
	default:
		return false
	}
}
