popiart-server/internal/server/store.go

package server

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"
)

type store struct {
	cfg        Config
	sessions   SessionRepository
	jobs       JobRepository
	routes     RouteRepository
	media      MediaRepository
	artifacts  ArtifactRepository
	skills     []skill
	projects   []project
	sessionTTL time.Duration
}

func newStore(cfg Config) (*store, error) {
	skills, err := loadSkills(cfg.SkillhubDir)
	if err != nil {
		return nil, err
	}
	projects := []project{
		{ID: "proj_local_dev", Name: "Local Dev"},
	}

	repo, err := newSQLiteRepository(cfg.SQLitePath, cfg.SessionSecret)
	if err != nil {
		return nil, err
	}

	return &store{
		cfg:        cfg,
		sessions:   repo,
		jobs:       repo,
		routes:     repo,
		media:      repo,
		artifacts:  repo,
		skills:     skills,
		projects:   projects,
		sessionTTL: 30 * 24 * time.Hour,
	}, nil
}

func seedSkills() []skill {
	return []skill{
		{
			ID:                 "popiskill-image-text2image-basic-v1",
			Name:               "Basic Text2Image",
			Description:        "Generate a single image from a text prompt through PopiNewAPI.",
			Tags:               []string{"image", "text2image", "basic"},
			Version:            "1.0.0",
			ModelType:          "image",
			RouteKey:           "image.text2image",
			EstimatedDurationS: 30,
			InputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"prompt": map[string]any{"type": "string"},
					"size":   map[string]any{"type": "string"},
				},
				"required": []string{"prompt"},
			},
			OutputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"artifact_ids": map[string]any{"type": "array", "items": map[string]any{"type": "string"}},
					"usage":        map[string]any{"type": "object"},
				},
			},
		},
		{
			ID:                 "popiskill-image-img2img-basic-v1",
			Name:               "Basic Img2Img",
			Description:        "Edit an existing image through PopiArt using a source artifact first, or a remote reference image URL as a fallback.",
			Tags:               []string{"image", "img2img", "basic"},
			Version:            "1.0.0",
			ModelType:          "image",
			RouteKey:           "image.img2img",
			EstimatedDurationS: 45,
			InputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"prompt":              map[string]any{"type": "string"},
					"source_artifact_id":  map[string]any{"type": "string"},
					"reference_image_url": map[string]any{"type": "string"},
					"image_url":           map[string]any{"type": "string"},
					"size":                map[string]any{"type": "string"},
					"aspect_ratio":        map[string]any{"type": "string"},
					"resolution":          map[string]any{"type": "string"},
					"notes":               map[string]any{"type": "string"},
				},
				"required": []string{"prompt"},
			},
			OutputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"artifact_ids": map[string]any{"type": "array", "items": map[string]any{"type": "string"}},
				},
			},
		},
		{
			ID:                 "popiskill-video-image2video-basic-v1",
			Name:               "Basic Image2Video",
			Description:        "Reserved image2video test skill. The runtime is not connected yet.",
			Tags:               []string{"video", "image2video", "basic"},
			Version:            "1.0.0",
			ModelType:          "video",
			RouteKey:           "video.image2video",
			EstimatedDurationS: 90,
			InputSchema:        imageToVideoInputSchema(),
			OutputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"artifact_ids": map[string]any{"type": "array", "items": map[string]any{"type": "string"}},
				},
			},
		},
	}
}

func (s *store) createSession(upstreamKey string) (string, user, bool, error) {
	current, err := s.sessions.CreateSession(upstreamKey, buildUserFromUpstreamKey(upstreamKey), s.sessionTTL)
	if err != nil {
		return "", user{}, false, err
	}
	return current.Token, current.User, true, nil
}

func (s *store) session(token string) (session, bool, error) {
	return s.sessions.GetSession(token)
}

func (s *store) deleteSession(token string) error {
	return s.sessions.DeleteSession(token)
}

func (s *store) rotateSession(oldToken string) (string, user, bool, error) {
	current, ok, err := s.sessions.RotateSession(oldToken, s.sessionTTL)
	if err != nil || !ok {
		return "", user{}, ok, err
	}
	return current.Token, current.User, true, nil
}

func (s *store) listSkills(tag, search string, limit, offset int) ([]skill, int) {
	filtered := make([]skill, 0, len(s.skills))
	for _, item := range s.skills {
		if tag != "" && !contains(item.Tags, tag) {
			continue
		}
		if search != "" {
			haystack := strings.ToLower(item.Name + " " + item.Description + " " + strings.Join(item.Tags, " "))
			if !strings.Contains(haystack, strings.ToLower(search)) {
				continue
			}
		}
		filtered = append(filtered, item)
	}

	total := len(filtered)
	start := clamp(offset, 0, total)
	end := clamp(offset+limit, start, total)
	return filtered[start:end], total
}

func (s *store) getSkill(id string) (skill, bool) {
	for _, item := range s.skills {
		if item.ID == id {
			return item, true
		}
	}
	return skill{}, false
}

func (s *store) createJob(skillID, routeKey, modelID, execMode string, input map[string]any, projectID, priority, idem string, current session) (*job, int, error) {
	now := time.Now().UTC()
	record := &job{
		JobID:          "job_" + randomID(),
		Status:         "pending",
		SkillID:        skillID,
		RouteKey:       normalizeRouteKey(routeKey),
		ModelID:        modelID,
		ExecMode:       defaultString(execMode, "sync_result"),
		Input:          cloneMap(input),
		ProjectID:      projectID,
		Priority:       priority,
		IdempotencyKey: idem,
		UserID:         current.User.ID,
		SessionID:      current.Token,
		CreatedAt:      now.Format(time.RFC3339),
		ArtifactIDs:    []string{},
		UpstreamKey:    current.UpstreamKey,
	}
	return s.jobs.CreateJob(record)
}

func (s *store) getJob(userID, jobID string) (*job, bool, error) {
	record, exists, err := s.jobs.GetJob(jobID)
	if err != nil || !exists {
		return nil, exists, err
	}
	if record.UserID != userID {
		return nil, false, nil
	}
	return record, true, nil
}

func (s *store) listJobs(userID, status, skillID, projectID string, limit, offset int) ([]job, int, error) {
	return s.jobs.ListJobs(userID, status, skillID, projectID, limit, offset)
}

func (s *store) cancelJob(userID, jobID string) (*job, bool, bool, error) {
	return s.jobs.CancelJob(userID, jobID)
}

func (s *store) jobLogs(userID, jobID string) ([]logEntry, bool, error) {
	record, exists, err := s.getJob(userID, jobID)
	if err != nil || !exists {
		return nil, exists, err
	}
	return buildJobLogs(record), true, nil
}

func (s *store) startJob(jobID string) (*job, bool, error) {
	return s.jobs.MarkJobRunning(jobID)
}

func (s *store) linkUpstreamTask(jobID, upstreamTaskID string) error {
	return s.jobs.LinkUpstreamTask(jobID, upstreamTaskID)
}

func (s *store) failJob(jobID, code, message string, details map[string]any) error {
	return s.jobs.FailJob(jobID, code, message, details)
}

func (s *store) completeJobWithResults(jobID string, refs []resultRef, usage map[string]any) error {
	if err := s.jobs.CompleteSyncResult(jobID, refs, cloneMap(usage)); err != nil {
		return err
	}
	record, exists, err := s.jobs.GetJob(jobID)
	if err != nil || !exists || record == nil {
		return err
	}
	return s.syncArtifactsForJob(record)
}

func (s *store) artifactsForJob(userID, jobID string) ([]artifact, bool, error) {
	record, exists, err := s.getJob(userID, jobID)
	if err != nil || !exists {
		return nil, exists, err
	}
	if err := s.syncArtifactsForJob(record); err != nil {
		return nil, true, err
	}
	if s.artifacts != nil {
		items, _, err := s.artifacts.ListArtifacts(userID, "", jobID, 1000, 0)
		if err != nil {
			return nil, true, err
		}
		if len(items) > 0 || len(record.ResultRefs) == 0 {
			return items, true, nil
		}
	}
	return buildArtifacts(record), true, nil
}

func (s *store) artifactRef(userID, artifactID string) (*artifact, resultRef, bool, error) {
	if s.artifacts != nil {
		item, ref, exists, err := s.artifacts.GetArtifact(userID, artifactID)
		if err != nil || exists {
			return item, ref, exists, err
		}
	}
	jobID, idx, ok := parseArtifactID(artifactID)
	if !ok {
		return nil, resultRef{}, false, nil
	}
	record, exists, err := s.getJob(userID, jobID)
	if err != nil || !exists {
		return nil, resultRef{}, exists, err
	}
	if err := s.syncArtifactsForJob(record); err != nil {
		return nil, resultRef{}, true, err
	}
	if s.artifacts != nil {
		item, ref, exists, err := s.artifacts.GetArtifact(userID, artifactID)
		if err != nil || exists {
			return item, ref, exists, err
		}
	}
	items := buildArtifacts(record)
	if idx < 0 || idx >= len(items) || idx >= len(record.ResultRefs) {
		return nil, resultRef{}, false, nil
	}
	item := items[idx]
	return &item, record.ResultRefs[idx], true, nil
}

func (s *store) listArtifacts(userID, projectID, jobID string, limit, offset int) ([]artifact, int, error) {
	if limit <= 0 {
		limit = 20
	}
	if offset < 0 {
		offset = 0
	}

	if err := s.backfillArtifacts(userID, strings.TrimSpace(projectID), strings.TrimSpace(jobID)); err != nil {
		return nil, 0, err
	}
	if s.artifacts != nil {
		return s.artifacts.ListArtifacts(userID, projectID, jobID, limit, offset)
	}
	return nil, 0, nil
}

func (s *store) routesForProject(projectID string) (map[string]string, error) {
	return s.routes.GetRoutes(projectID)
}

func (s *store) setRouteOverride(projectID, routeKey, modelID string) error {
	return s.routes.SetRoute(projectID, normalizeRouteKey(routeKey), modelID)
}

func (s *store) unsetRouteOverride(projectID, routeKey string) error {
	return s.routes.UnsetRoute(projectID, normalizeRouteKey(routeKey))
}

func (s *store) persistMediaContent(userID, projectID, artifactID, filename, contentType, visibility string, content []byte, mediaID string) (mediaRecord, error) {
	return persistMediaContent(s.cfg, s.media, userID, projectID, artifactID, filename, contentType, visibility, content, mediaID)
}

func (s *store) getMedia(mediaID string) (*mediaRecord, bool, error) {
	if s.media != nil {
		record, exists, err := s.media.GetMedia(mediaID)
		if err != nil || exists {
			return record, exists, err
		}
	}
	record, exists, err := loadMediaRecordFromJSON(s.cfg, mediaID)
	if err != nil || !exists || record == nil {
		return record, exists, err
	}
	if s.media != nil {
		if upsertErr := s.media.UpsertMedia(*record); upsertErr != nil {
			return nil, false, upsertErr
		}
	}
	return record, true, nil
}

func (s *store) syncArtifactsForJob(record *job) error {
	if s.artifacts == nil || record == nil {
		return nil
	}
	return s.artifacts.UpsertArtifacts(record.UserID, record.ProjectID, buildArtifacts(record))
}

func (s *store) backfillArtifacts(userID, projectID, jobID string) error {
	if s.artifacts == nil {
		return nil
	}
	if jobID != "" {
		record, exists, err := s.getJob(userID, jobID)
		if err != nil || !exists || record == nil {
			return err
		}
		if projectID != "" && record.ProjectID != projectID {
			return nil
		}
		return s.syncArtifactsForJob(record)
	}

	const jobPageSize = 200
	totalJobs := 0
	for jobOffset := 0; ; jobOffset += jobPageSize {
		jobs, total, err := s.jobs.ListJobs(userID, "", "", projectID, jobPageSize, jobOffset)
		if err != nil {
			return err
		}
		totalJobs = total
		for idx := range jobs {
			if err := s.syncArtifactsForJob(&jobs[idx]); err != nil {
				return err
			}
		}
		if jobOffset+len(jobs) >= totalJobs || len(jobs) == 0 {
			break
		}
	}
	return nil
}

func cloneJob(record *job) *job {
	if record == nil {
		return nil
	}
	data := *record
	data.Input = cloneMap(record.Input)
	if record.ResultRefs != nil {
		data.ResultRefs = append([]resultRef{}, record.ResultRefs...)
	}
	if record.ArtifactIDs != nil {
		data.ArtifactIDs = append([]string{}, record.ArtifactIDs...)
	}
	return &data
}

func buildUserFromUpstreamKey(key string) user {
	sum := sha256.Sum256([]byte(key))
	fingerprint := hex.EncodeToString(sum[:6])
	short := fingerprint
	if len(short) > 12 {
		short = short[:12]
	}
	return user{
		ID:     "user_" + short,
		Email:  "",
		Name:   "PopiNewAPI Session",
		Scopes: []string{"skills:read", "jobs:write", "artifacts:read", "artifacts:write", "budget:read"},
	}
}

func buildJobLogs(record *job) []logEntry {
	if record == nil {
		return nil
	}
	logs := []logEntry{
		{TS: record.CreatedAt, Level: "info", Message: "job accepted"},
	}
	if record.StartedAt != "" {
		message := "job dispatched to PopiNewAPI"
		if record.ExecMode == "upstream_task" && record.NewAPITaskID != "" {
			message = fmt.Sprintf("job linked to PopiNewAPI task %s", record.NewAPITaskID)
		}
		logs = append(logs, logEntry{TS: record.StartedAt, Level: "info", Message: message})
	}
	switch record.Status {
	case "done":
		logs = append(logs, logEntry{TS: defaultString(record.FinishedAt, record.CreatedAt), Level: "info", Message: "job completed"})
	case "failed":
		message := "job failed"
		if record.Error != nil && strings.TrimSpace(record.Error.Message) != "" {
			message = record.Error.Message
		}
		logs = append(logs, logEntry{TS: defaultString(record.FinishedAt, record.CreatedAt), Level: "error", Message: message})
	case "cancelled":
		logs = append(logs, logEntry{TS: defaultString(record.FinishedAt, record.CreatedAt), Level: "warn", Message: "job cancelled"})
	}
	return logs
}

func buildArtifacts(record *job) []artifact {
	if record == nil || len(record.ResultRefs) == 0 {
		return nil
	}
	createdAt := defaultString(record.FinishedAt, record.CreatedAt)
	expiresAt := ""
	if created := parseRFC3339(createdAt); !created.IsZero() {
		expiresAt = created.Add(30 * 24 * time.Hour).Format(time.RFC3339)
	}
	items := make([]artifact, 0, len(record.ResultRefs))
	for idx, ref := range record.ResultRefs {
		filename := strings.TrimSpace(ref.Filename)
		if filename == "" {
			filename = inferFilenameFromRef(ref, buildArtifactID(record.JobID, idx))
		}
		contentType := defaultString(strings.TrimSpace(ref.ContentType), "application/octet-stream")
		storageStatus := strings.TrimSpace(ref.StorageStatus)
		if storageStatus == "" {
			switch {
			case ref.LocalPath != "":
				storageStatus = "ready"
			case ref.URL != "":
				storageStatus = "upstream"
			case ref.DataURL != "":
				storageStatus = "embedded"
			}
		}
		item := artifact{
			ID:             buildArtifactID(record.JobID, idx),
			JobID:          record.JobID,
			ProjectID:      record.ProjectID,
			MediaID:        strings.TrimSpace(ref.MediaID),
			Filename:       filename,
			ContentType:    contentType,
			SizeBytes:      ref.SizeBytes,
			CreatedAt:      createdAt,
			ExpiresAt:      expiresAt,
			URL:            strings.TrimSpace(ref.URL),
			Visibility:     strings.TrimSpace(ref.Visibility),
			SHA256:         strings.TrimSpace(ref.SHA256),
			StorageStatus:  storageStatus,
			SourceSkillID:  strings.TrimSpace(record.SkillID),
			SourceModelID:  strings.TrimSpace(record.ModelID),
			SourceRouteKey: strings.TrimSpace(record.RouteKey),
			SourceInput:    cloneMap(record.Input),
			PromptText:     primaryPromptText(record.Input),
			Usage:          cloneMap(record.Usage),
			Ref:            ref,
		}
		items = append(items, item)
	}
	return items
}

func buildArtifactIDs(jobID string, refs []resultRef) []string {
	if len(refs) == 0 {
		return nil
	}
	ids := make([]string, 0, len(refs))
	for idx := range refs {
		ids = append(ids, buildArtifactID(jobID, idx))
	}
	return ids
}

func primaryPromptText(input map[string]any) string {
	if input == nil {
		return ""
	}
	return strings.TrimSpace(stringValue(
		input["prompt"],
		input["motion_prompt"],
		input["scene_prompt"],
		input["text"],
	))
}

func buildArtifactID(jobID string, idx int) string {
	return "art_" + strings.TrimPrefix(jobID, "job_") + "_" + strconv.Itoa(idx)
}

func parseArtifactID(artifactID string) (string, int, bool) {
	if !strings.HasPrefix(artifactID, "art_") {
		return "", 0, false
	}
	value := strings.TrimPrefix(artifactID, "art_")
	pos := strings.LastIndex(value, "_")
	if pos <= 0 {
		return "", 0, false
	}
	idx, err := strconv.Atoi(value[pos+1:])
	if err != nil {
		return "", 0, false
	}
	return "job_" + value[:pos], idx, true
}

func inferFilenameFromRef(ref resultRef, fallback string) string {
	if ref.URL != "" {
		if parsed, err := url.Parse(ref.URL); err == nil {
			base := path.Base(parsed.Path)
			if base != "" && base != "." && base != "/" {
				return base
			}
		}
	}
	return fallback
}

func normalizeRouteKey(value string) string {
	switch strings.TrimSpace(value) {
	case "", "image":
		return "image.text2image"
	case "video":
		return "video.image2video"
	case "audio":
		return "audio.default"
	default:
		return value
	}
}

func routeExecMode(routeKey string) string {
	switch normalizeRouteKey(routeKey) {
	case "video.image2video":
		return "upstream_task"
	default:
		return "sync_result"
	}
}

func maskSecret(key string) string {
	if key == "" {
		return ""
	}
	if len(key) <= 12 {
		return key
	}
	return key[:4] + "********" + key[len(key)-4:]
}

func contains(items []string, target string) bool {
	for _, item := range items {
		if strings.EqualFold(item, target) {
			return true
		}
	}
	return false
}

func clamp(v, min, max int) int {
	if v < min {
		return min
	}
	if v > max {
		return max
	}
	return v
}

func cloneMap(src map[string]any) map[string]any {
	if src == nil {
		return nil
	}
	dst := make(map[string]any, len(src))
	for key, value := range src {
		dst[key] = value
	}
	return dst
}

func randomID() string {
	buf := make([]byte, 6)
	if _, err := rand.Read(buf); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return hex.EncodeToString(buf)
}