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"
	"syscall"
	"time"

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

const (
	localBinaryResultDirName          = "results"
	localBinaryPlaceholderPrefix      = "[GatewayBinary:v1;"
	localBinaryGenericBase64MinLength = 4096
	localBinaryMaxDepth               = 64
	defaultLocalResultTTLHours        = 24
	defaultLocalResultMinFreeBytes    = int64(10 * 1024 * 1024 * 1024)
	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
	taskDir       string
	writeFiles    bool
	enforceLimits bool
	totalBytes    int64
	seen          map[string]struct{}
	createdFiles  []string
}

// materializeLocalBinaryResult replaces every inline binary value with a
// bounded placeholder after atomically writing and verifying the bytes locally.
func (s *Service) materializeLocalBinaryResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
	if _, _, err := s.transformLocalBinaryResult(ctx, taskID, result, false, true); err != nil {
		return nil, false, err
	}
	return s.transformLocalBinaryResult(ctx, taskID, result, true, true)
}

func (s *Service) transformLocalBinaryResult(ctx context.Context, taskID string, result map[string]any, writeFiles bool, enforceLimits bool) (map[string]any, bool, error) {
	root := s.localBinaryResultRoot()
	taskDir := filepath.Join(root, safeLocalBinaryTaskDir(taskID))
	materializer := &localBinaryMaterializer{
		service:       s,
		taskDir:       taskDir,
		writeFiles:    writeFiles,
		enforceLimits: enforceLimits,
		seen:          map[string]struct{}{},
	}
	next, changed, err := materializer.materializeValue(ctx, result, "", nil, 0)
	if err != nil {
		materializer.rollbackCreatedFiles()
		return nil, false, err
	}
	mapped, ok := next.(map[string]any)
	if !ok {
		materializer.rollbackCreatedFiles()
		return nil, false, &clients.ClientError{
			Code:       "result_binary_not_materialized",
			Message:    "generated result is not a JSON object",
			StatusCode: 500,
			Retryable:  false,
		}
	}
	if changed && enforceLimits && !writeFiles {
		if err := os.MkdirAll(root, 0o750); err != nil {
			return nil, false, localBinaryStorageError(err)
		}
		if err := ensureLocalBinaryDiskHeadroom(root, materializer.totalBytes, s.localResultMinFreeBytes()); err != nil {
			return nil, false, err
		}
	}
	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) {
	return s.materializeLocalBinaryResult(ctx, taskID, result)
}

// CompactExpiredTaskResultForStorage computes the same deterministic
// placeholders without creating files that are already outside the recovery
// window.
func (s *Service) CompactExpiredTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
	return s.transformLocalBinaryResult(ctx, taskID, result, false, false)
}

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,
			}
		}
		if m.writeFiles {
			created, err := m.service.writeLocalBinaryResult(m.taskDir, digestHex, payload)
			if err != nil {
				return nil, false, err
			}
			if created {
				m.createdFiles = append(m.createdFiles, filepath.Join(m.taskDir, digestHex+".bin"))
			}
		}
		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 (m *localBinaryMaterializer) rollbackCreatedFiles() {
	for _, path := range m.createdFiles {
		_ = os.Remove(path)
	}
	_ = os.Remove(m.taskDir)
	m.createdFiles = nil
}

func (s *Service) writeLocalBinaryResult(taskDir string, digestHex string, payload []byte) (bool, error) {
	if err := os.MkdirAll(taskDir, 0o750); err != nil {
		return false, localBinaryStorageError(err)
	}
	targetPath := filepath.Join(taskDir, digestHex+".bin")
	if info, err := os.Stat(targetPath); err == nil {
		if !info.IsDir() && info.Size() == int64(len(payload)) {
			if err := verifyLocalBinaryFile(targetPath, digestHex, int64(len(payload))); err == nil {
				now := time.Now()
				if err := os.Chtimes(targetPath, now, now); err != nil {
					return false, localBinaryStorageError(err)
				}
				return false, nil
			}
		}
		return false, &clients.ClientError{
			Code:       "binary_result_corrupted",
			Message:    "existing local result file does not match its content hash",
			StatusCode: 500,
			Retryable:  false,
		}
	} else if !errors.Is(err, os.ErrNotExist) {
		return false, localBinaryStorageError(err)
	}
	if err := ensureLocalBinaryDiskHeadroom(taskDir, int64(len(payload)), s.localResultMinFreeBytes()); err != nil {
		return false, err
	}
	tempFile, err := os.CreateTemp(taskDir, ".gateway-result-*")
	if err != nil {
		return false, localBinaryStorageError(err)
	}
	tempPath := tempFile.Name()
	cleanup := func() {
		_ = tempFile.Close()
		_ = os.Remove(tempPath)
	}
	if err := tempFile.Chmod(0o640); err != nil {
		cleanup()
		return false, localBinaryStorageError(err)
	}
	if _, err := tempFile.Write(payload); err != nil {
		cleanup()
		return false, localBinaryStorageError(err)
	}
	if err := tempFile.Sync(); err != nil {
		cleanup()
		return false, localBinaryStorageError(err)
	}
	if err := tempFile.Close(); err != nil {
		_ = os.Remove(tempPath)
		return false, localBinaryStorageError(err)
	}
	if err := os.Rename(tempPath, targetPath); err != nil {
		_ = os.Remove(tempPath)
		return false, localBinaryStorageError(err)
	}
	if err := verifyLocalBinaryFile(targetPath, digestHex, int64(len(payload))); err != nil {
		_ = os.Remove(targetPath)
		return false, err
	}
	return true, nil
}

func ensureLocalBinaryDiskHeadroom(path string, incomingBytes int64, minFreeBytes int64) error {
	var stat syscall.Statfs_t
	if err := syscall.Statfs(path, &stat); err != nil {
		return localBinaryStorageError(err)
	}
	freeBytes := int64(stat.Bavail) * int64(stat.Bsize)
	if freeBytes-incomingBytes < minFreeBytes {
		return &clients.ClientError{
			Code:       "local_result_storage_unavailable",
			Message:    "local result storage does not have enough free space",
			StatusCode: 503,
			Retryable:  false,
		}
	}
	return 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:
		next := make(map[string]any, len(typed))
		changed := false
		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 (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) localResultMinFreeBytes() int64 {
	if s.cfg.LocalResultMinFreeBytes <= 0 {
		return defaultLocalResultMinFreeBytes
	}
	return s.cfg.LocalResultMinFreeBytes
}

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
}