package windsurf

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"strings"
)

// CascadeImage 是发给 Windsurf Cascade gRPC 的图像载体。
// 对应 proto：message CodeiumImage { string base64Data = 1; string mimeType = 2; string caption = 3; }
// 通过静态分析 /Applications/Windsurf.app/Contents/Resources/app/node_modules/@exa/chat-client/index.js 得到。
type CascadeImage struct {
	// Base64Data 原始 base64 字符串（不含 data: 前缀）。仅用于 replay/发送；不参与指纹。
	Base64Data string `json:"base64_data,omitempty"`
	// MimeType 例如 image/png。参与指纹。
	MimeType string `json:"mime_type"`
	// Caption 可选辅助说明，参与指纹。
	Caption string `json:"caption,omitempty"`
	// SHA256 是 decoded 二进制字节的 sha256 hex；指纹使用。
	SHA256 string `json:"sha256,omitempty"`
	// ByteLen decoded 字节数；指纹使用。
	ByteLen int `json:"byte_len,omitempty"`
}

// ImageDigest 是 CascadeImage 的摘要视图。日志/指纹/conversation pool 使用这个，
// 永远不要把 Base64Data 带进哈希/持久化。
type ImageDigest struct {
	MimeType string `json:"mime_type"`
	SHA256   string `json:"sha256"`
	ByteLen  int    `json:"byte_len"`
	Caption  string `json:"caption,omitempty"`
}

// 支持的图像 MIME（与 Windsurf 客户端一致）
var cascadeImageAllowedMime = map[string]struct{}{
	"image/png":     {},
	"image/jpeg":    {},
	"image/gif":     {},
	"image/svg+xml": {},
}

// CascadeImageMaxPerTurn 单次 user turn 最多图片数（Windsurf 客户端限制）。
const CascadeImageMaxPerTurn = 5

// CascadeImageMaxBytes 单张解码后字节上限（Windsurf 客户端压缩目标 1MB）。
const CascadeImageMaxBytes = 1 * 1024 * 1024

// CascadeImageValidationOptions 校验开关，便于不同场景选择严格度。
type CascadeImageValidationOptions struct {
	// EnforceByteSize 是否对单张做 1MB 上限校验。默认 true。
	EnforceByteSize bool
	// EnforceCountPerTurn 是否对张数做 <=5 校验。默认 true。
	EnforceCountPerTurn bool
}

// DefaultCascadeImageValidationOptions 返回默认校验选项。
func DefaultCascadeImageValidationOptions() CascadeImageValidationOptions {
	return CascadeImageValidationOptions{
		EnforceByteSize:     true,
		EnforceCountPerTurn: true,
	}
}

// ValidateCascadeImages 在发送前做确定性校验。
// 返回第一条校验失败的错误，上层据此生成 Anthropic 风格 400。
func ValidateCascadeImages(images []CascadeImage, opts CascadeImageValidationOptions) error {
	if opts.EnforceCountPerTurn && len(images) > CascadeImageMaxPerTurn {
		return fmt.Errorf("too many images: %d (max %d)", len(images), CascadeImageMaxPerTurn)
	}
	for i := range images {
		img := &images[i]
		if _, ok := cascadeImageAllowedMime[strings.ToLower(strings.TrimSpace(img.MimeType))]; !ok {
			return fmt.Errorf("image[%d]: unsupported media_type %q", i, img.MimeType)
		}
		trimmed := strings.TrimSpace(img.Base64Data)
		if trimmed == "" {
			return fmt.Errorf("image[%d]: data must not be empty", i)
		}
		decoded, err := base64.StdEncoding.DecodeString(trimmed)
		if err != nil {
			// 尝试 RawStdEncoding 作为兜底（部分客户端省略 padding）
			if decoded2, err2 := base64.RawStdEncoding.DecodeString(trimmed); err2 == nil {
				decoded = decoded2
			} else {
				return fmt.Errorf("image[%d]: invalid base64 data", i)
			}
		}
		if len(decoded) == 0 {
			return fmt.Errorf("image[%d]: decoded data is empty", i)
		}
		if opts.EnforceByteSize && len(decoded) > CascadeImageMaxBytes {
			return fmt.Errorf("image[%d]: decoded size %d exceeds 1MB limit", i, len(decoded))
		}
		// 顺手把 digest/byteLen 算好，后续指纹阶段直接用
		if img.SHA256 == "" || img.ByteLen == 0 {
			sum := sha256.Sum256(decoded)
			img.SHA256 = hex.EncodeToString(sum[:])
			img.ByteLen = len(decoded)
		}
		// 归一化 MIME（避免大小写差异）
		img.MimeType = strings.ToLower(strings.TrimSpace(img.MimeType))
	}
	return nil
}

// BuildImageDigests 把 CascadeImage 转成只含摘要的 ImageDigest 切片。
// 前提：调用者已经通过 ValidateCascadeImages 或 ComputeImageDigest 确保 SHA256/ByteLen 已填。
func BuildImageDigests(images []CascadeImage) []ImageDigest {
	if len(images) == 0 {
		return nil
	}
	out := make([]ImageDigest, len(images))
	for i, img := range images {
		out[i] = ImageDigest{
			MimeType: img.MimeType,
			SHA256:   img.SHA256,
			ByteLen:  img.ByteLen,
			Caption:  img.Caption,
		}
	}
	return out
}

// ComputeImageDigest 按给定 base64 字符串计算 digest 字段（不做限额校验）。
// 给只需要摘要不需要发送场景使用（例如历史 replay 构造）。
func ComputeImageDigest(img *CascadeImage) error {
	decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(img.Base64Data))
	if err != nil {
		decoded, err = base64.RawStdEncoding.DecodeString(strings.TrimSpace(img.Base64Data))
		if err != nil {
			return fmt.Errorf("invalid base64")
		}
	}
	sum := sha256.Sum256(decoded)
	img.SHA256 = hex.EncodeToString(sum[:])
	img.ByteLen = len(decoded)
	img.MimeType = strings.ToLower(strings.TrimSpace(img.MimeType))
	return nil
}

// encodeCodeiumImage 对单张 CodeiumImage 消息进行 proto wire 编码：
// message CodeiumImage { string base64Data = 1; string mimeType = 2; string caption = 3; }
// 返回值是 message body 原始字节（不含外层 field tag/length）。
func encodeCodeiumImage(img CascadeImage) []byte {
	var body []byte
	// field 1：base64 字符串原样发（Windsurf 客户端也是原样放，不做二次编码）
	body = append(body, encodeStringField(1, img.Base64Data)...)
	body = append(body, encodeStringField(2, img.MimeType)...)
	if img.Caption != "" {
		body = append(body, encodeStringField(3, img.Caption)...)
	}
	return body
}

// appendSendUserCascadeImages 把一组图像追加到 SendUserCascadeMessageRequest.images（field 6, repeated）。
// repeated message 字段按"每条独立的 field 6 segment"编码，与 Windsurf 客户端保持一致。
func appendSendUserCascadeImages(body []byte, images []CascadeImage) []byte {
	for _, img := range images {
		body = append(body, encodeBytesField(6, encodeCodeiumImage(img))...)
	}
	return body
}