Luminary/internal/prediction/service.go

package prediction

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/provider/comfyui"
	"github.com/rose_cat707/luminary/internal/storage/imagestore"
	"github.com/rose_cat707/luminary/internal/store/cache"
	"gorm.io/gorm"
)

type Service struct {
	db        *gorm.DB
	cache     *cache.Store
	images    *imagestore.Store
	hub     *Hub
	waitSec int64
	mu        sync.Mutex
	running   map[string]context.CancelFunc
}

func NewService(db *gorm.DB, c *cache.Store, images *imagestore.Store, waitSec int64) *Service {
	if waitSec <= 0 {
		waitSec = 60
	}
	return &Service{
		db: db, cache: c, images: images, hub: NewHub(),
		waitSec: waitSec,
		running: make(map[string]context.CancelFunc),
	}
}

func (s *Service) Hub() *Hub { return s.hub }

func NewPredictionID() (string, error) {
	b := make([]byte, 12)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

type CreateRequest struct {
	Version string         `json:"version"`
	Input   map[string]any `json:"input"`
}

func (s *Service) Create(ctx context.Context, ingress *model.IngressKey, req CreateRequest, clientIP string, wait bool) (*model.Prediction, error) {
	if req.Version == "" {
		return nil, fmt.Errorf("version is required")
	}
	entry, provider, err := s.resolveModel(req.Version, ingress)
	if err != nil {
		return nil, err
	}
	if !provider.Enabled {
		return nil, fmt.Errorf("provider disabled")
	}
	input, err := ValidateInputWithDefaults(entry.WorkflowType, req.Input, inputDefaultsForEntry(entry))
	if err != nil {
		return nil, err
	}
	binding, err := ParseBinding(entry.InputBindingJSON)
	if err != nil {
		return nil, err
	}
	var workflow map[string]any
	_ = json.Unmarshal([]byte(entry.WorkflowJSON), &workflow)
	if err := binding.Validate(workflow, RequiredBindings(entry.WorkflowType)); err != nil {
		return nil, fmt.Errorf("model binding invalid: %w", err)
	}

	id, err := NewPredictionID()
	if err != nil {
		return nil, err
	}
	ingressModelID := entry.IngressID()
	inputJSON, _ := json.Marshal(input)
	p := model.Prediction{
		ID:           id,
		IngressKeyID: ingress.ID,
		ProviderID:   provider.ID,
		ModelEntryID: entry.ID,
		Version:      ingressModelID,
		Model:        ingressModelID,
		InputJSON:    string(inputJSON),
		Status:       model.PredictionStarting,
		ClientID:     "luminary-" + id,
		CreatedAt:    time.Now(),
	}
	if err := s.db.Create(&p).Error; err != nil {
		return nil, err
	}

	runCtx, cancel := context.WithCancel(context.Background())
	s.mu.Lock()
	s.running[id] = cancel
	s.mu.Unlock()

	go s.execute(runCtx, &p, entry, &provider, binding, input, clientIP)

	if wait {
		deadline := time.After(time.Duration(s.waitSec) * time.Second)
		for {
			select {
			case <-deadline:
				return s.Get(ctx, id, ingress)
			case <-time.After(200 * time.Millisecond):
				cur, err := s.Get(ctx, id, ingress)
				if err != nil {
					return nil, err
				}
				if cur.Status == model.PredictionSucceeded || cur.Status == model.PredictionFailed || cur.Status == model.PredictionCanceled {
					return cur, nil
				}
			}
		}
	}
	return s.toAPI(&p), nil
}

func inputDefaultsForEntry(entry model.ModelEntry) map[string]any {
	defs, err := ParseInputDefaults(entry.InputDefaultsJSON)
	if err != nil {
		return map[string]any{}
	}
	return defs
}

func (s *Service) resolveModel(version string, ingress *model.IngressKey) (model.ModelEntry, model.Provider, error) {
	allowedProviders := cache.ParseJSONList(ingress.AllowedProvidersJSON)
	allowedModels := cache.ParseJSONList(ingress.AllowedModelsJSON)
	var candidates []model.ModelEntry
	for _, p := range s.cache.ListProviders() {
		if p.Category != model.CategoryImage || !p.Enabled {
			continue
		}
		if !cache.MatchesList(allowedProviders, p.Name) && !cache.MatchesList(allowedProviders, string(p.Type)) && !cache.MatchesList(allowedProviders, "*") {
			continue
		}
		for _, m := range s.cache.GetModels(p.ID) {
			if !m.Enabled {
				continue
			}
			ingressID := m.IngressID()
			if !m.MatchesIngressRef(version, p.Name) {
				continue
			}
			if !cache.MatchesList(allowedModels, ingressID) && !cache.MatchesList(allowedModels, m.ModelID) && !cache.MatchesList(allowedModels, "*") {
				continue
			}
			candidates = append(candidates, m)
		}
	}
	if len(candidates) == 0 {
		return model.ModelEntry{}, model.Provider{}, fmt.Errorf("no model found for version %s", version)
	}
	if len(candidates) > 1 && !strings.Contains(version, "/") {
		return model.ModelEntry{}, model.Provider{}, fmt.Errorf("ambiguous model name %q; use owner/name", version)
	}
	if len(candidates) > 1 {
		return model.ModelEntry{}, model.Provider{}, fmt.Errorf("ambiguous model %q", version)
	}
	m := candidates[0]
	p, ok := s.cache.GetProvider(m.ProviderID)
	if !ok {
		return model.ModelEntry{}, model.Provider{}, fmt.Errorf("provider not found")
	}
	return m, p, nil
}

func (s *Service) execute(ctx context.Context, p *model.Prediction, entry model.ModelEntry, provider *model.Provider, binding InputBinding, input map[string]any, clientIP string) {
	defer func() {
		s.mu.Lock()
		delete(s.running, p.ID)
		s.mu.Unlock()
	}()

	start := time.Now()
	s.updateStatus(p, model.PredictionProcessing, "", nil)
	s.hub.Publish(p.ID, Event{Status: string(model.PredictionProcessing)})

	client := comfyui.New(provider.BaseURL)
	runInput := cloneMap(input)
	if entry.WorkflowType == model.WorkflowImg2Img {
		imageURL, _ := runInput["image"].(string)
		data, name, err := s.loadInputImage(ctx, imageURL)
		if err != nil {
			s.fail(p, clientIP, start, fmt.Errorf("download image: %w", err))
			return
		}
		uploaded, err := client.UploadImage(ctx, data, name)
		if err != nil {
			s.fail(p, clientIP, start, fmt.Errorf("upload image: %w", err))
			return
		}
		runInput["image"] = uploaded
	}

	prompt, err := BuildPrompt(entry.WorkflowJSON, binding, runInput)
	if err != nil {
		s.fail(p, clientIP, start, err)
		return
	}

	wsCtx, wsCancel := context.WithCancel(ctx)
	defer wsCancel()
	go client.Watch(wsCtx, p.ClientID, func(ev comfyui.ProgressEvent) {
		switch ev.Type {
		case "progress":
			if ev.Max > 0 {
				logs := fmt.Sprintf("progress %.0f/%.0f", ev.Value, ev.Max)
				s.appendLogs(p, logs)
				s.hub.Publish(p.ID, Event{Status: string(model.PredictionProcessing), Logs: logs})
			}
		case "execution_error":
			s.appendLogs(p, ev.ErrorText)
		}
	})

	submit, err := client.SubmitPrompt(ctx, prompt, p.ClientID)
	if err != nil {
		s.fail(p, clientIP, start, err)
		return
	}
	p.ComfyPromptID = submit.PromptID
	now := time.Now()
	p.StartedAt = &now
	s.db.Model(p).Updates(map[string]any{"comfy_prompt_id": p.ComfyPromptID, "started_at": p.StartedAt})

	for {
		select {
		case <-ctx.Done():
			_ = client.Interrupt(context.Background())
			s.updateStatus(p, model.PredictionCanceled, "canceled", nil)
			s.hub.Publish(p.ID, Event{Status: string(model.PredictionCanceled), Done: true})
			s.recordUsage(ingressByID(s.cache, p.IngressKeyID), provider, clientIP, p, start, "canceled", "canceled")
			return
		case <-time.After(2 * time.Second):
			history, err := client.GetHistory(ctx, p.ComfyPromptID)
			if err != nil {
				continue
			}
			if _, ok := history[p.ComfyPromptID]; !ok {
				continue
			}
			images := comfyui.CollectOutputImages(history, p.ComfyPromptID)
			if len(images) == 0 {
				continue
			}
			var urls []string
			for _, img := range images {
				data, mime, err := client.DownloadView(ctx, img.Filename, img.Subfolder, img.Type)
				if err != nil {
					s.fail(p, clientIP, start, err)
					return
				}
				stored, err := s.images.Save(p.ID, img.Filename, data, mime)
				if err != nil {
					s.fail(p, clientIP, start, err)
					return
				}
				urls = append(urls, s.images.SignURL(stored.ID))
			}
			outJSON, _ := json.Marshal(urls)
			p.OutputJSON = string(outJSON)
			completed := time.Now()
			p.CompletedAt = &completed
			metrics, _ := json.Marshal(map[string]any{
				"predict_time": time.Since(start).Seconds(),
				"total_time":   completed.Sub(p.CreatedAt).Seconds(),
			})
			p.MetricsJSON = string(metrics)
			s.updateStatus(p, model.PredictionSucceeded, "", &completed)
			s.hub.Publish(p.ID, Event{
				Status: string(model.PredictionSucceeded), Output: urls,
				Metrics: map[string]any{"predict_time": time.Since(start).Seconds()},
				Done: true,
			})
			s.recordUsage(ingressByID(s.cache, p.IngressKeyID), provider, clientIP, p, start, "success", "")
			return
		}
	}
}

func (s *Service) fail(p *model.Prediction, clientIP string, start time.Time, err error) {
	completed := time.Now()
	p.CompletedAt = &completed
	p.Error = err.Error()
	s.updateStatus(p, model.PredictionFailed, err.Error(), &completed)
	s.hub.Publish(p.ID, Event{Status: string(model.PredictionFailed), Error: err.Error(), Done: true})
	prov, _ := s.cache.GetProvider(p.ProviderID)
	s.recordUsage(ingressByID(s.cache, p.IngressKeyID), &prov, clientIP, p, start, "error", err.Error())
}

func (s *Service) updateStatus(p *model.Prediction, status model.PredictionStatus, errMsg string, completed *time.Time) {
	p.Status = status
	if errMsg != "" {
		p.Error = errMsg
	}
	updates := map[string]any{"status": status, "error": p.Error, "logs": p.Logs, "output_json": p.OutputJSON, "metrics_json": p.MetricsJSON}
	if completed != nil {
		updates["completed_at"] = completed
	}
	s.db.Model(p).Where("id = ?", p.ID).Updates(updates)
}

func (s *Service) appendLogs(p *model.Prediction, line string) {
	if p.Logs != "" {
		p.Logs += "\n"
	}
	p.Logs += line
	s.db.Model(p).Where("id = ?", p.ID).Update("logs", p.Logs)
}

func (s *Service) Get(ctx context.Context, id string, ingress *model.IngressKey) (*model.Prediction, error) {
	var p model.Prediction
	if err := s.db.First(&p, "id = ?", id).Error; err != nil {
		return nil, fmt.Errorf("prediction not found")
	}
	if ingress != nil && p.IngressKeyID != ingress.ID {
		return nil, fmt.Errorf("prediction not found")
	}
	return s.toAPI(&p), nil
}

func (s *Service) List(ctx context.Context, ingress *model.IngressKey, cursor string, limit int) ([]*model.Prediction, error) {
	if limit <= 0 || limit > 100 {
		limit = 100
	}
	q := s.db.Where("ingress_key_id = ?", ingress.ID).Order("created_at desc").Limit(limit)
	if cursor != "" {
		q = q.Where("id < ?", cursor)
	}
	var rows []model.Prediction
	if err := q.Find(&rows).Error; err != nil {
		return nil, err
	}
	out := make([]*model.Prediction, 0, len(rows))
	for i := range rows {
		out = append(out, s.toAPI(&rows[i]))
	}
	return out, nil
}

func (s *Service) Cancel(ctx context.Context, id string, ingress *model.IngressKey) (*model.Prediction, error) {
	p, err := s.Get(ctx, id, ingress)
	if err != nil {
		return nil, err
	}
	if p.Status == model.PredictionSucceeded || p.Status == model.PredictionFailed || p.Status == model.PredictionCanceled {
		return p, nil
	}
	s.mu.Lock()
	if cancel, ok := s.running[id]; ok {
		cancel()
	}
	s.mu.Unlock()
	return s.Get(ctx, id, ingress)
}

func (s *Service) toAPI(p *model.Prediction) *model.Prediction {
	return p
}

func (s *Service) APIView(p *model.Prediction, baseURL string) map[string]any {
	var input any
	_ = json.Unmarshal([]byte(p.InputJSON), &input)
	var output any
	if p.OutputJSON != "" {
		_ = json.Unmarshal([]byte(p.OutputJSON), &output)
	}
	var metrics any
	if p.MetricsJSON != "" {
		_ = json.Unmarshal([]byte(p.MetricsJSON), &metrics)
	}
	base := strings.TrimRight(baseURL, "/")
	return map[string]any{
		"id":           p.ID,
		"version":      p.Version,
		"input":        input,
		"output":       output,
		"error":        nilIfEmpty(p.Error),
		"status":       p.Status,
		"created_at":   p.CreatedAt.UTC().Format(time.RFC3339Nano),
		"started_at":   formatTime(p.StartedAt),
		"completed_at": formatTime(p.CompletedAt),
		"logs":         p.Logs,
		"metrics":      metrics,
		"urls": map[string]string{
			"get":    fmt.Sprintf("%s/v1/predictions/%s", base, p.ID),
			"cancel": fmt.Sprintf("%s/v1/predictions/%s/cancel", base, p.ID),
			"stream": fmt.Sprintf("%s/v1/predictions/%s/stream", base, p.ID),
		},
	}
}

func nilIfEmpty(s string) any {
	if s == "" {
		return nil
	}
	return s
}

func formatTime(t *time.Time) any {
	if t == nil {
		return nil
	}
	return t.UTC().Format(time.RFC3339Nano)
}

func (s *Service) loadInputImage(ctx context.Context, imageURL string) ([]byte, string, error) {
	if id, ok := parseGatewayFileID(imageURL); ok && s.images != nil {
		data, img, err := s.images.ReadContent(id)
		if err == nil {
			if err := validateImageBytes(data); err != nil {
				return nil, "", err
			}
			return data, imageFilenameFromStored(img.Filename, img.Mime, data), nil
		}
	}
	data, err := downloadURL(ctx, imageURL)
	if err != nil {
		return nil, "", err
	}
	if err := validateImageBytes(data); err != nil {
		return nil, "", err
	}
	return data, imageFilenameFromURL(imageURL, data), nil
}

func downloadURL(ctx context.Context, raw string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, raw, nil)
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("http %d", resp.StatusCode)
	}
	return io.ReadAll(resp.Body)
}

func cloneMap(in map[string]any) map[string]any {
	b, _ := json.Marshal(in)
	var out map[string]any
	_ = json.Unmarshal(b, &out)
	return out
}

func ingressByID(c *cache.Store, id uint) *model.IngressKey {
	for _, k := range c.ListIngressKeys() {
		if k.ID == id {
			kCopy := k
			return &kCopy
		}
	}
	return nil
}

func (s *Service) recordUsage(ingress *model.IngressKey, provider *model.Provider, clientIP string, p *model.Prediction, start time.Time, status, errMsg string) {
	if ingress == nil {
		return
	}
	var providerID uint
	if provider != nil {
		providerID = provider.ID
	}
	s.cache.RecordUsage(cache.UsageDelta{
		IngressKeyID: ingress.ID,
		ProviderID:   providerID,
		ClientIP:     clientIP,
		Endpoint:     model.EndpointPrediction,
		Model:        p.Model,
		LatencyMs:    time.Since(start).Milliseconds(),
		Status:       status,
		ErrorMsg:     errMsg,
		RequestBody:  p.InputJSON,
		ResponseBody: p.OutputJSON,
		Requests:     1,
	})
}

func (s *Service) GetDBPrediction(id string) (*model.Prediction, error) {
	var p model.Prediction
	if err := s.db.First(&p, "id = ?", id).Error; err != nil {
		return nil, err
	}
	return &p, nil
}