easyai-ai-gateway/apps/api/internal/runner/binary_results.go

package runner

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

const (
	localBinaryResultDirName          = "results"
	localBinaryPlaceholderPrefix      = "[GatewayBinary:v1;"
	localBinaryGenericBase64MinLength = 4096
	localBinaryMaxDepth               = 64
	defaultLocalResultTTLHours        = 24
	defaultLocalResultMaxBytes        = int64(256 * 1024 * 1024)
	defaultLocalResultMaxTaskBytes    = int64(512 * 1024 * 1024)
)

type localBinaryDescriptor struct {
	Prefix      string
	SHA256      string
	Size        int64
	ContentType string
	Encoding    string
}

type localBinaryMaterializer struct {
	service       *Service
	enforceLimits bool
	totalBytes    int64
	seen          map[string]struct{}
}

// materializeLocalBinaryResult is retained as an internal compatibility seam,
// but new materialization always targets object storage. Historical local
// placeholders remain readable through HydrateTaskResult.
func (s *Service) materializeLocalBinaryResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
	hadInline := TaskResultHasInlineBinary(result)
	if !hadInline {
		return result, false, nil
	}
	next, err := s.uploadGeneratedAssets(ctx, taskID, "", "", result)
	if err != nil {
		return nil, false, err
	}
	return next, !TaskResultHasInlineBinary(next), nil
}

func (s *Service) transformLocalBinaryResult(ctx context.Context, _ string, result map[string]any, _ bool, enforceLimits bool) (map[string]any, bool, error) {
	materializer := &localBinaryMaterializer{
		service:       s,
		enforceLimits: enforceLimits,
		seen:          map[string]struct{}{},
	}
	next, changed, err := materializer.materializeValue(ctx, result, "", nil, 0)
	if err != nil {
		return nil, false, err
	}
	mapped, ok := next.(map[string]any)
	if !ok {
		return nil, false, &clients.ClientError{
			Code:       "result_binary_not_materialized",
			Message:    "generated result is not a JSON object",
			StatusCode: 500,
			Retryable:  false,
		}
	}
	return mapped, changed, nil
}

// MaterializeTaskResultForStorage exposes the same verified materialization
// path to the explicit historical maintenance command.
func (s *Service) MaterializeTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
	hadInline := TaskResultHasInlineBinary(result)
	if !hadInline {
		return result, false, nil
	}
	next, err := s.uploadGeneratedAssets(ctx, taskID, "", "", result)
	if err != nil {
		return nil, false, err
	}
	return next, hadInline && !TaskResultHasInlineBinary(next), nil
}

// CompactExpiredTaskResultForStorage keeps the historical maintenance API but
// now uses the same object-storage path as live results. New GatewayBinary
// placeholders are never created; their parser remains read-only compatibility.
func (s *Service) CompactExpiredTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
	return s.MaterializeTaskResultForStorage(ctx, taskID, result)
}

func TaskResultHasInlineBinary(result map[string]any) bool {
	return localBinaryValueHasPayload(result, "", nil, 0)
}

func localBinaryValueHasPayload(value any, key string, siblings map[string]any, depth int) bool {
	if depth >= localBinaryMaxDepth {
		return false
	}
	switch typed := value.(type) {
	case map[string]any:
		if _, _, ok := localBufferObjectBytes(typed); ok {
			return true
		}
		for childKey, child := range typed {
			if localBinaryValueHasPayload(child, childKey, typed, depth+1) {
				return true
			}
		}
	case []any:
		if localBinaryKey(key) {
			if _, ok := bytesFromNumberArray(typed); ok {
				return true
			}
		}
		for _, child := range typed {
			if localBinaryValueHasPayload(child, key, siblings, depth+1) {
				return true
			}
		}
	case []byte:
		return len(typed) > 0
	case string:
		_, _, _, ok := localBinaryStringBytes(key, typed, siblings)
		return ok
	}
	return false
}

func (m *localBinaryMaterializer) materializeValue(ctx context.Context, value any, key string, siblings map[string]any, depth int) (any, bool, error) {
	if err := ctx.Err(); err != nil {
		return nil, false, err
	}
	if depth >= localBinaryMaxDepth {
		return nil, false, &clients.ClientError{
			Code:       "result_binary_not_materialized",
			Message:    "generated result exceeds the maximum JSON depth",
			StatusCode: 500,
			Retryable:  false,
		}
	}
	switch typed := value.(type) {
	case map[string]any:
		if payload, contentType, ok := localBufferObjectBytes(typed); ok {
			return m.persistBinary(ctx, payload, contentType, "buffer")
		}
		next := make(map[string]any, len(typed))
		changed := false
		for childKey, childValue := range typed {
			child, childChanged, err := m.materializeValue(ctx, childValue, childKey, typed, depth+1)
			if err != nil {
				return nil, false, err
			}
			next[childKey] = child
			changed = changed || childChanged
		}
		if !changed {
			return value, false, nil
		}
		return next, true, nil
	case []any:
		if localBinaryKey(key) {
			if payload, ok := bytesFromNumberArray(typed); ok {
				return m.persistBinary(ctx, payload, mediaContentTypeFromItem(siblings), "buffer")
			}
		}
		next := make([]any, len(typed))
		changed := false
		for index, item := range typed {
			child, childChanged, err := m.materializeValue(ctx, item, key, siblings, depth+1)
			if err != nil {
				return nil, false, err
			}
			next[index] = child
			changed = changed || childChanged
		}
		if !changed {
			return value, false, nil
		}
		return next, true, nil
	case []byte:
		if len(typed) == 0 {
			return value, false, nil
		}
		return m.persistBinary(ctx, append([]byte(nil), typed...), mediaContentTypeFromItem(siblings), "buffer")
	case string:
		payload, contentType, encoding, ok := localBinaryStringBytes(key, typed, siblings)
		if !ok {
			return value, false, nil
		}
		return m.persistBinary(ctx, payload, contentType, encoding)
	default:
		return value, false, nil
	}
}

func (m *localBinaryMaterializer) persistBinary(ctx context.Context, payload []byte, contentType string, encoding string) (any, bool, error) {
	if err := ctx.Err(); err != nil {
		return nil, false, err
	}
	if len(payload) == 0 {
		return nil, false, nil
	}
	size := int64(len(payload))
	if m.enforceLimits && size > m.service.localResultMaxBytes() {
		return nil, false, &clients.ClientError{
			Code:       "binary_result_too_large",
			Message:    "one generated binary result exceeds the local storage limit",
			StatusCode: 502,
			Retryable:  false,
		}
	}
	digest := sha256.Sum256(payload)
	digestHex := hex.EncodeToString(digest[:])
	if _, exists := m.seen[digestHex]; !exists {
		if m.enforceLimits && m.totalBytes+size > m.service.localResultMaxTaskBytes() {
			return nil, false, &clients.ClientError{
				Code:       "binary_result_too_large",
				Message:    "generated binary results exceed the per-task local storage limit",
				StatusCode: 502,
				Retryable:  false,
			}
		}
		m.seen[digestHex] = struct{}{}
		m.totalBytes += size
	}
	descriptor := localBinaryDescriptor{
		Prefix:      localBinaryPrefix(payload),
		SHA256:      digestHex,
		Size:        size,
		ContentType: normalizedLocalBinaryContentType(contentType),
		Encoding:    normalizedLocalBinaryEncoding(encoding),
	}
	return localBinaryPlaceholder(descriptor), true, nil
}

func localBinaryStorageError(err error) error {
	return &clients.ClientError{
		Code:       "local_result_storage_unavailable",
		Message:    "local result storage failed: " + err.Error(),
		StatusCode: 503,
		Retryable:  false,
	}
}

func verifyLocalBinaryFile(path string, expectedHash string, expectedSize int64) error {
	file, err := os.Open(path)
	if err != nil {
		return localBinaryStorageError(err)
	}
	defer file.Close()
	hasher := sha256.New()
	size, err := io.Copy(hasher, file)
	if err != nil {
		return localBinaryStorageError(err)
	}
	if size != expectedSize || hex.EncodeToString(hasher.Sum(nil)) != expectedHash {
		return &clients.ClientError{
			Code:       "binary_result_corrupted",
			Message:    "local result file failed size or hash verification",
			StatusCode: 500,
			Retryable:  false,
		}
	}
	return nil
}

// HydrateTaskResult restores placeholders from verified local files. It is only
// used by result/detail/replay endpoints, never by task lists or callbacks.
func (s *Service) HydrateTaskResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, error) {
	next, changed, err := s.hydrateLocalBinaryValue(ctx, taskID, result, 0)
	if err != nil {
		return nil, err
	}
	if !changed {
		return result, nil
	}
	mapped, ok := next.(map[string]any)
	if !ok {
		return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "stored result is not a JSON object", StatusCode: 500}
	}
	return mapped, nil
}

func (s *Service) hydrateLocalBinaryValue(ctx context.Context, taskID string, value any, depth int) (any, bool, error) {
	if err := ctx.Err(); err != nil {
		return nil, false, err
	}
	if depth >= localBinaryMaxDepth {
		return nil, false, &clients.ClientError{Code: "binary_result_corrupted", Message: "stored result exceeds the maximum JSON depth", StatusCode: 500}
	}
	switch typed := value.(type) {
	case map[string]any:
		refreshedAccessURL := false
		if asset, upload, ok := generatedResultUploadReference(typed); ok {
			accessURL, err := s.requestAssetAccessURL(ctx, asset)
			if err != nil {
				return nil, false, err
			}
			next := make(map[string]any, len(typed))
			for key, item := range typed {
				next[key] = item
			}
			nextUpload := make(map[string]any, len(upload))
			for key, item := range upload {
				nextUpload[key] = item
			}
			nextUpload["url"] = accessURL
			next["upload"] = nextUpload
			for _, key := range []string{"url", "image_url", "video_url", "audio_url"} {
				if _, exists := next[key]; exists {
					next[key] = accessURL
				}
			}
			typed = next
			refreshedAccessURL = true
		}
		if ref, ok := generatedResultAssetReference(typed); ok {
			payload, contentType, err := s.readGeneratedResultAsset(ctx, ref)
			if err != nil {
				return nil, false, err
			}
			encoded := base64.StdEncoding.EncodeToString(payload)
			if generatedResultAssetUsesDataURI(typed) {
				return "data:" + contentType + ";base64," + encoded, true, nil
			}
			return encoded, true, nil
		}
		next := make(map[string]any, len(typed))
		changed := refreshedAccessURL
		for key, childValue := range typed {
			child, childChanged, err := s.hydrateLocalBinaryValue(ctx, taskID, childValue, depth+1)
			if err != nil {
				return nil, false, err
			}
			next[key] = child
			changed = changed || childChanged
		}
		if !changed {
			return value, false, nil
		}
		return next, true, nil
	case []any:
		next := make([]any, len(typed))
		changed := false
		for index, childValue := range typed {
			child, childChanged, err := s.hydrateLocalBinaryValue(ctx, taskID, childValue, depth+1)
			if err != nil {
				return nil, false, err
			}
			next[index] = child
			changed = changed || childChanged
		}
		if !changed {
			return value, false, nil
		}
		return next, true, nil
	case string:
		descriptor, ok := parseLocalBinaryPlaceholder(typed)
		if !ok {
			return value, false, nil
		}
		payload, err := s.readLocalBinaryResult(taskID, descriptor)
		if err != nil {
			return nil, false, err
		}
		encoded := base64.StdEncoding.EncodeToString(payload)
		if descriptor.Encoding == "data-uri" {
			return "data:" + descriptor.ContentType + ";base64," + encoded, true, nil
		}
		return encoded, true, nil
	default:
		return value, false, nil
	}
}

func generatedResultUploadReference(value map[string]any) (store.RequestAsset, map[string]any, bool) {
	upload, ok := value["upload"].(map[string]any)
	if !ok {
		return store.RequestAsset{}, nil, false
	}
	objectKey := strings.TrimSpace(stringFromAny(upload["objectKey"]))
	channel, _ := upload["storageChannel"].(map[string]any)
	channelKey := strings.TrimSpace(stringFromAny(channel["channelKey"]))
	if objectKey == "" || channelKey == "" {
		return store.RequestAsset{}, nil, false
	}
	return store.RequestAsset{
		URL:               stringFromAny(upload["url"]),
		StorageProvider:   stringFromAny(channel["provider"]),
		StorageChannelID:  stringFromAny(channel["id"]),
		StorageChannelKey: channelKey,
		ObjectKey:         objectKey,
		AccessScope:       stringFromAny(upload["accessScope"]),
	}, upload, true
}

func generatedResultAssetReference(value map[string]any) (store.RequestAsset, bool) {
	ref, ok := value["assetRef"].(map[string]any)
	if !ok {
		return store.RequestAsset{}, false
	}
	storage, _ := value["assetStorage"].(map[string]any)
	if stringFromAny(storage["scene"]) != store.FileStorageSceneImageResult {
		return store.RequestAsset{}, false
	}
	asset := store.RequestAsset{
		SHA256:            strings.ToLower(strings.TrimSpace(stringFromAny(ref["sha256"]))),
		ContentType:       firstNonEmptyString(stringFromAny(ref["contentType"]), stringFromAny(storage["contentType"])),
		URL:               firstNonEmptyString(stringFromAny(ref["url"]), stringFromAny(value["url"])),
		StorageProvider:   stringFromAny(ref["storageProvider"]),
		StorageChannelID:  stringFromAny(ref["storageChannelId"]),
		StorageChannelKey: stringFromAny(ref["storageChannelKey"]),
		ObjectKey:         stringFromAny(ref["objectKey"]),
		AccessScope:       stringFromAny(ref["accessScope"]),
	}
	if size := floatFromAny(ref["size"]); size > 0 {
		asset.ByteSize = int64(size)
	}
	if expiresAt := stringFromAny(ref["expiresAt"]); expiresAt != "" {
		if parsed, err := time.Parse(time.RFC3339, expiresAt); err == nil {
			asset.ExpiresAt = &parsed
		}
	}
	if asset.URL == "" || asset.SHA256 == "" || asset.ByteSize <= 0 {
		return store.RequestAsset{}, false
	}
	return asset, true
}

func generatedResultAssetUsesDataURI(value map[string]any) bool {
	storage, _ := value["assetStorage"].(map[string]any)
	source := normalizeLocalBinaryKey(stringFromAny(storage["source"]))
	return source == "datauri"
}

func (s *Service) readGeneratedResultAsset(ctx context.Context, asset store.RequestAsset) ([]byte, string, error) {
	payload, err := s.readRequestAssetBytes(ctx, asset)
	if err != nil {
		return nil, "", err
	}
	digest := sha256.Sum256(payload)
	if int64(len(payload)) != asset.ByteSize || hex.EncodeToString(digest[:]) != asset.SHA256 {
		return nil, "", &clients.ClientError{
			Code:       "binary_result_corrupted",
			Message:    "stored result asset failed size or hash verification",
			StatusCode: 500,
			Retryable:  false,
		}
	}
	contentType := strings.TrimSpace(asset.ContentType)
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	return payload, contentType, nil
}

func (s *Service) readLocalBinaryResult(taskID string, descriptor localBinaryDescriptor) ([]byte, error) {
	path := filepath.Join(s.localBinaryResultRoot(), safeLocalBinaryTaskDir(taskID), descriptor.SHA256+".bin")
	info, err := os.Stat(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, &clients.ClientError{Code: "binary_result_expired", Message: "local binary result has expired", StatusCode: 410, Retryable: false}
		}
		return nil, localBinaryStorageError(err)
	}
	if info.IsDir() || info.Size() != descriptor.Size {
		return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "local binary result has an invalid size", StatusCode: 500, Retryable: false}
	}
	if info.ModTime().Before(time.Now().Add(-time.Duration(s.localResultTTLHours()) * time.Hour)) {
		return nil, &clients.ClientError{Code: "binary_result_expired", Message: "local binary result has expired", StatusCode: 410, Retryable: false}
	}
	payload, err := os.ReadFile(path)
	if err != nil {
		return nil, localBinaryStorageError(err)
	}
	digest := sha256.Sum256(payload)
	if int64(len(payload)) != descriptor.Size || hex.EncodeToString(digest[:]) != descriptor.SHA256 {
		return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "local binary result failed size or hash verification", StatusCode: 500, Retryable: false}
	}
	return payload, nil
}

func localBinaryStringBytes(key string, value string, siblings map[string]any) ([]byte, string, string, bool) {
	raw := strings.TrimSpace(value)
	if raw == "" || strings.HasPrefix(raw, localBinaryPlaceholderPrefix) {
		return nil, "", "", false
	}
	if strings.HasPrefix(strings.ToLower(raw), "data:") {
		contentType, encoded, ok, err := parseBase64DataURL(raw)
		if err == nil && ok {
			payload, decodeErr := decodeBase64Payload(encoded)
			if decodeErr == nil && len(payload) > 0 {
				return payload, contentType, "data-uri", true
			}
		}
		return nil, "", "", false
	}
	strict := localBinaryKey(key)
	if !strict && len(raw) < localBinaryGenericBase64MinLength {
		return nil, "", "", false
	}
	payload, err := decodeBase64Payload(raw)
	if err != nil || len(payload) == 0 {
		return nil, "", "", false
	}
	return payload, firstNonEmptyString(mediaContentTypeFromItem(siblings), defaultContentTypeForRawMediaKey(key)), "raw", true
}

func localBufferObjectBytes(value map[string]any) ([]byte, string, bool) {
	if normalizeLocalBinaryKey(stringFromAny(value["type"])) != "buffer" {
		return nil, "", false
	}
	contentType := firstNonEmptyString(
		stringFromAny(value["mime_type"]),
		stringFromAny(value["mimeType"]),
		stringFromAny(value["contentType"]),
	)
	switch data := value["data"].(type) {
	case []byte:
		if len(data) == 0 {
			return nil, "", false
		}
		return append([]byte(nil), data...), contentType, true
	case []any:
		payload, ok := bytesFromNumberArray(data)
		return payload, contentType, ok
	default:
		return nil, "", false
	}
}

func localBinaryKey(key string) bool {
	normalized := normalizeLocalBinaryKey(key)
	return normalized == "b64" ||
		normalized == "b64json" ||
		normalized == "base64" ||
		normalized == "buffer" ||
		normalized == "bytes" ||
		strings.Contains(normalized, "base64") ||
		strings.Contains(normalized, "buffer") ||
		strings.Contains(normalized, "binary") ||
		strings.HasSuffix(normalized, "b64") ||
		strings.HasSuffix(normalized, "bytes")
}

func normalizeLocalBinaryKey(value string) string {
	return strings.Map(func(char rune) rune {
		switch {
		case char >= 'a' && char <= 'z':
			return char
		case char >= 'A' && char <= 'Z':
			return char + ('a' - 'A')
		case char >= '0' && char <= '9':
			return char
		default:
			return -1
		}
	}, value)
}

func localBinaryPrefix(payload []byte) string {
	prefix := base64.StdEncoding.EncodeToString(payload)
	if len(prefix) > 16 {
		prefix = prefix[:16]
	}
	return prefix
}

func normalizedLocalBinaryContentType(value string) string {
	value = normalizeGeneratedContentType(value)
	if value == "" || len(value) > 32 || !localBinaryContentTypeSafe(value) {
		return "application/octet-stream"
	}
	return value
}

func localBinaryContentTypeSafe(value string) bool {
	for _, char := range value {
		switch {
		case char >= 'a' && char <= 'z':
		case char >= 'A' && char <= 'Z':
		case char >= '0' && char <= '9':
		case char == '/', char == '.', char == '+', char == '-':
		default:
			return false
		}
	}
	return true
}

func normalizedLocalBinaryEncoding(value string) string {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "data-uri":
		return "data-uri"
	case "buffer":
		return "buffer"
	default:
		return "raw"
	}
}

func localBinaryPlaceholder(descriptor localBinaryDescriptor) string {
	return fmt.Sprintf(
		"[GatewayBinary:v1;prefix=%s;sha256=%s;bytes=%d;mime=%s;encoding=%s]",
		descriptor.Prefix,
		descriptor.SHA256,
		descriptor.Size,
		descriptor.ContentType,
		descriptor.Encoding,
	)
}

func parseLocalBinaryPlaceholder(value string) (localBinaryDescriptor, bool) {
	if !strings.HasPrefix(value, localBinaryPlaceholderPrefix) || !strings.HasSuffix(value, "]") {
		return localBinaryDescriptor{}, false
	}
	content := strings.TrimSuffix(strings.TrimPrefix(value, localBinaryPlaceholderPrefix), "]")
	fields := map[string]string{}
	for _, item := range strings.Split(content, ";") {
		key, fieldValue, ok := strings.Cut(item, "=")
		if !ok {
			continue
		}
		fields[key] = fieldValue
	}
	size, err := strconv.ParseInt(fields["bytes"], 10, 64)
	if err != nil || size <= 0 || len(fields["sha256"]) != sha256.Size*2 {
		return localBinaryDescriptor{}, false
	}
	if _, err := hex.DecodeString(fields["sha256"]); err != nil {
		return localBinaryDescriptor{}, false
	}
	descriptor := localBinaryDescriptor{
		Prefix:      fields["prefix"],
		SHA256:      strings.ToLower(fields["sha256"]),
		Size:        size,
		ContentType: normalizedLocalBinaryContentType(fields["mime"]),
		Encoding:    normalizedLocalBinaryEncoding(fields["encoding"]),
	}
	if len(descriptor.Prefix) > 16 {
		return localBinaryDescriptor{}, false
	}
	return descriptor, true
}

func localBinaryResultHasPlaceholders(value any) bool {
	switch typed := value.(type) {
	case map[string]any:
		for _, child := range typed {
			if localBinaryResultHasPlaceholders(child) {
				return true
			}
		}
	case []any:
		for _, child := range typed {
			if localBinaryResultHasPlaceholders(child) {
				return true
			}
		}
	case string:
		return strings.HasPrefix(typed, localBinaryPlaceholderPrefix)
	}
	return false
}

func safeLocalBinaryTaskDir(taskID string) string {
	taskID = strings.TrimSpace(taskID)
	if taskID != "" {
		safe := true
		for _, char := range taskID {
			if (char >= 'a' && char <= 'z') ||
				(char >= 'A' && char <= 'Z') ||
				(char >= '0' && char <= '9') ||
				char == '-' || char == '_' {
				continue
			}
			safe = false
			break
		}
		if safe && taskID != "." && taskID != ".." && len(taskID) <= 128 {
			return taskID
		}
	}
	digest := sha256.Sum256([]byte(taskID))
	return "task-" + hex.EncodeToString(digest[:16])
}

func (s *Service) localBinaryResultRoot() string {
	root := strings.TrimSpace(s.cfg.LocalGeneratedStorageDir)
	if root == "" {
		root = config.DefaultLocalGeneratedStorageDir
	}
	return filepath.Join(root, localBinaryResultDirName)
}

func (s *Service) localResultTTLHours() int {
	if s.cfg.LocalResultTTLHours <= 0 {
		return defaultLocalResultTTLHours
	}
	return s.cfg.LocalResultTTLHours
}

func (s *Service) localResultMaxBytes() int64 {
	if s.cfg.LocalResultMaxBytes <= 0 {
		return defaultLocalResultMaxBytes
	}
	return s.cfg.LocalResultMaxBytes
}

func (s *Service) localResultMaxTaskBytes() int64 {
	if s.cfg.LocalResultMaxTaskBytes <= 0 {
		return defaultLocalResultMaxTaskBytes
	}
	return s.cfg.LocalResultMaxTaskBytes
}