preface-tools/internal/tools/comicanimator/worker.go

package comicanimator

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"git.michelsen.id/phill/preface-tools/internal/tools/comicanimator/openrouter"
)

func (t *Tool) Run(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return nil
		case gid := <-t.queue:
			func() {
				defer func() {
					if v := recover(); v != nil {
						t.fail(gid, "internal_error", "Generation failed unexpectedly.", fmt.Errorf("panic: %v", v))
					}
				}()
				t.work(ctx, gid)
			}()
		}
	}
}
func (t *Tool) work(appctx context.Context, id string) {
	g, ok := t.store.generation(id)
	if !ok {
		return
	}
	t.store.update(id, func(x *Generation) { x.Status = Submitting })
	ctx, cancel := context.WithTimeout(appctx, t.cfg.JobTimeout)
	defer cancel()
	job, err := t.client.Submit(ctx, openrouter.VideoRequest{Model: t.cfg.VideoModel, Prompt: g.ReviewedPrompt, Duration: g.Duration, Resolution: g.Resolution, AspectRatio: g.AspectRatio, GenerateAudio: t.cfg.GenerateAudio, FrameImages: []openrouter.FrameImage{{Type: "image_url", ImageURL: openrouter.NewImageURL(t.signedURL(g.UploadID)), FrameType: "first_frame"}}})
	if err != nil {
		t.fail(id, "provider_submission_failed", "The video provider could not start this generation.", err)
		return
	}
	if job.ID == "" {
		t.fail(id, "provider_submission_failed", "The video provider returned an invalid response.", fmt.Errorf("missing job id"))
		return
	}
	t.store.update(id, func(x *Generation) { x.ProviderJobID = job.ID; x.Status = Pending })
	for job.Status != "completed" {
		if job.Status == "failed" || job.Status == "cancelled" || job.Status == "expired" {
			t.fail(id, "provider_generation_failed", "The video provider could not complete this generation.", fmt.Errorf("provider status %s", job.Status))
			return
		}
		timer := time.NewTimer(t.cfg.PollInterval)
		select {
		case <-ctx.Done():
			timer.Stop()
			code := "provider_timeout"
			msg := "The video generation timed out."
			if appctx.Err() != nil {
				code = "worker_shutdown"
				msg = "Generation stopped because the application shut down."
			}
			t.fail(id, code, msg, ctx.Err())
			return
		case <-timer.C:
		}
		job, err = t.client.Poll(ctx, job.ID)
		if err != nil {
			t.fail(id, "provider_poll_failed", "Could not check the video status.", err)
			return
		}
		t.store.update(id, func(x *Generation) {
			if job.Status == "in_progress" {
				x.Status = InProgress
			} else {
				x.Status = Pending
			}
		})
	}
	t.store.update(id, func(x *Generation) { x.Status = Downloading })
	resp, err := t.client.Content(ctx, job.ID)
	if err != nil {
		t.fail(id, "provider_download_failed", "Could not download the completed video.", err)
		return
	}
	defer resp.Body.Close()
	mt := strings.Split(resp.Header.Get("Content-Type"), ";")[0]
	if mt != "video/mp4" && mt != "video/webm" && mt != "video/quicktime" {
		t.fail(id, "generated_video_invalid", "The provider returned an invalid video.", fmt.Errorf("content type %q", mt))
		return
	}
	ext := map[string]string{"video/mp4": ".mp4", "video/webm": ".webm", "video/quicktime": ".mov"}[mt]
	header := make([]byte, 16)
	hn, readErr := io.ReadFull(resp.Body, header)
	if readErr != nil && readErr != io.ErrUnexpectedEOF {
		t.fail(id, "provider_download_failed", "Could not download the completed video.", readErr)
		return
	}
	if !validVideoHeader(mt, header[:hn]) {
		t.fail(id, "generated_video_invalid", "The provider returned an invalid video.", fmt.Errorf("invalid %s signature", mt))
		return
	}
	tmp, err := os.CreateTemp(t.cfg.OutputDirectory, ".video-*")
	if err != nil {
		t.fail(id, "output_write_failed", "Could not save the video.", err)
		return
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName)
	n, copyErr := io.Copy(tmp, io.LimitReader(io.MultiReader(bytes.NewReader(header[:hn]), resp.Body), t.cfg.MaxGeneratedVideoBytes+1))
	syncErr := tmp.Sync()
	closeErr := tmp.Close()
	if copyErr != nil || syncErr != nil || closeErr != nil {
		t.fail(id, "output_write_failed", "Could not save the video.", errorsJoin(copyErr, syncErr, closeErr))
		return
	}
	if n > t.cfg.MaxGeneratedVideoBytes {
		t.fail(id, "generated_video_too_large", "The generated video was too large.", fmt.Errorf("%d bytes", n))
		return
	}
	now := time.Now().UTC()
	name := now.Format("20060102T150405Z") + "_" + id + ext
	dest := filepath.Join(t.cfg.OutputDirectory, name)
	if err = os.Rename(tmpName, dest); err != nil {
		t.fail(id, "output_write_failed", "Could not finalize the video.", err)
		return
	}
	t.store.update(id, func(x *Generation) {
		x.Status = Completed
		x.OutputPath = dest
		x.OutputMIMEType = mt
		x.OutputSize = n
		x.CompletedAt = &now
	})
	t.log.Info("video generation completed", "tool", "comic-animator", "generation_id", id, "provider_job_id", job.ID, "file_size", n)
}

func validVideoHeader(mt string, head []byte) bool {
	switch mt {
	case "video/mp4", "video/quicktime":
		return len(head) >= 12 && string(head[4:8]) == "ftyp"
	case "video/webm":
		return len(head) >= 4 && head[0] == 0x1a && head[1] == 0x45 && head[2] == 0xdf && head[3] == 0xa3
	default:
		return false
	}
}
func errorsJoin(errs ...error) error {
	parts := []string{}
	for _, e := range errs {
		if e != nil {
			parts = append(parts, e.Error())
		}
	}
	if len(parts) == 0 {
		return nil
	}
	return fmt.Errorf("%s", strings.Join(parts, "; "))
}
func (t *Tool) fail(id, code, message string, err error) {
	t.store.update(id, func(x *Generation) { x.Status = Failed; x.ErrorCode = code; x.ErrorMessage = message })
	t.log.Error("video generation failed", "tool", "comic-animator", "generation_id", id, "error_code", code, "error", err)
}