popiart-server/internal/server/media.go

package server

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type mediaRecord struct {
	media
	UserID    string `json:"user_id"`
	LocalPath string `json:"local_path"`
}

func mediaDir(cfg Config) string {
	return filepath.Join(cfg.DataDir, "media")
}

func mediaBlobDir(cfg Config) string {
	return filepath.Join(mediaDir(cfg), "blobs")
}

func mediaMetaDir(cfg Config) string {
	return filepath.Join(mediaDir(cfg), "meta")
}

func mediaMetaPath(cfg Config, mediaID string) string {
	return filepath.Join(mediaMetaDir(cfg), mediaID+".json")
}

func mediaBlobPath(cfg Config, mediaID, filename string) string {
	ext := filepath.Ext(filename)
	if ext == "" {
		ext = extensionFromContentType("")
	}
	return filepath.Join(mediaBlobDir(cfg), mediaID+ext)
}

func publicBaseURL(cfg Config) string {
	base := strings.TrimSpace(cfg.PublicBaseURL)
	if base == "" {
		base = strings.TrimSpace(os.Getenv("POPIART_SERVER_ADDR"))
	}
	if base == "" {
		base = "127.0.0.1:8080"
	}
	if !strings.HasPrefix(base, "http://") && !strings.HasPrefix(base, "https://") {
		base = "http://" + base
	}
	base = strings.Replace(base, "0.0.0.0", "127.0.0.1", 1)
	return strings.TrimRight(base, "/")
}

func mediaContentURL(cfg Config, mediaID string) string {
	return publicBaseURL(cfg) + "/v1/media/" + url.PathEscape(strings.TrimSpace(mediaID)) + "/content"
}

func persistMediaContent(cfg Config, userID, projectID, artifactID, filename, contentType, visibility string, content []byte, mediaID string) (mediaRecord, error) {
	if len(content) == 0 {
		return mediaRecord{}, fmt.Errorf("media content is required")
	}
	filename = sanitizeFilename(strings.TrimSpace(filename))
	if filename == "" || filename == "artifact" {
		filename = "artifact" + extensionFromContentType(contentType)
	}
	contentType = strings.TrimSpace(contentType)
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	visibility = strings.TrimSpace(visibility)
	if visibility == "" {
		visibility = "unlisted"
	}
	mediaID = strings.TrimSpace(mediaID)
	if mediaID == "" {
		mediaID = "med_" + randomID()
	}
	if err := os.MkdirAll(mediaBlobDir(cfg), 0o755); err != nil {
		return mediaRecord{}, err
	}
	if err := os.MkdirAll(mediaMetaDir(cfg), 0o755); err != nil {
		return mediaRecord{}, err
	}

	sum := sha256.Sum256(content)
	blobPath := mediaBlobPath(cfg, mediaID, filename)
	if err := os.WriteFile(blobPath, content, 0o644); err != nil {
		return mediaRecord{}, err
	}

	record := mediaRecord{
		media: media{
			ID:          mediaID,
			ArtifactID:  strings.TrimSpace(artifactID),
			ProjectID:   strings.TrimSpace(projectID),
			Filename:    filename,
			ContentType: contentType,
			SizeBytes:   int64(len(content)),
			CreatedAt:   time.Now().UTC().Format(time.RFC3339),
			URL:         mediaContentURL(cfg, mediaID),
			Visibility:  visibility,
			SHA256:      hex.EncodeToString(sum[:]),
		},
		UserID:    strings.TrimSpace(userID),
		LocalPath: blobPath,
	}

	metaBytes, err := json.Marshal(record)
	if err != nil {
		return mediaRecord{}, err
	}
	if err := os.WriteFile(mediaMetaPath(cfg, mediaID), metaBytes, 0o644); err != nil {
		return mediaRecord{}, err
	}
	return record, nil
}

func loadMediaRecord(cfg Config, mediaID string) (*mediaRecord, bool, error) {
	mediaID = strings.TrimSpace(mediaID)
	if mediaID == "" {
		return nil, false, nil
	}
	data, err := os.ReadFile(mediaMetaPath(cfg, mediaID))
	if os.IsNotExist(err) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	var record mediaRecord
	if err := json.Unmarshal(data, &record); err != nil {
		return nil, false, err
	}
	return &record, true, nil
}

func (s *Server) persistResultRefs(ctx context.Context, record *job, refs []resultRef) ([]resultRef, error) {
	if record == nil || len(refs) == 0 {
		return refs, nil
	}
	items := make([]resultRef, 0, len(refs))
	for idx, ref := range refs {
		contentType, _, reader, err := s.newapi.openResultRef(ctx, record.UpstreamKey, ref)
		if err != nil {
			return nil, err
		}
		content, readErr := io.ReadAll(reader)
		reader.Close()
		if readErr != nil {
			return nil, readErr
		}
		filename := strings.TrimSpace(ref.Filename)
		if filename == "" {
			filename = inferFilenameFromRef(ref, buildArtifactID(record.JobID, idx))
		}
		artifactID := buildArtifactID(record.JobID, idx)
		mediaRecord, err := persistMediaContent(
			s.cfg,
			record.UserID,
			record.ProjectID,
			artifactID,
			filename,
			defaultString(strings.TrimSpace(contentType), strings.TrimSpace(ref.ContentType)),
			defaultString(strings.TrimSpace(ref.Visibility), "unlisted"),
			content,
			"",
		)
		if err != nil {
			return nil, err
		}
		items = append(items, resultRef{
			Kind:          "local_path",
			URL:           mediaRecord.URL,
			LocalPath:     mediaRecord.LocalPath,
			MediaID:       mediaRecord.ID,
			Filename:      mediaRecord.Filename,
			ContentType:   mediaRecord.ContentType,
			SizeBytes:     mediaRecord.SizeBytes,
			Visibility:    mediaRecord.Visibility,
			SHA256:        mediaRecord.SHA256,
			StorageStatus: "ready",
		})
	}
	return items, nil
}