llink/go/cmd/particleprocessorworker/main.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"log/slog"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/flowy-live/llink/internal/billing"
	"github.com/flowy-live/llink/internal/db"
	"github.com/flowy-live/llink/internal/depot"
	"github.com/flowy-live/llink/internal/particle"
	"github.com/flowy-live/llink/internal/speech"
	"github.com/flowy-live/llink/internal/utils"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"cloud.google.com/go/firestore"
	"cloud.google.com/go/storage"
)

func createClient(ctx context.Context) *firestore.Client {
	projectId := utils.MustGetEnv("GCP_PROJECT")

	client, err := firestore.NewClient(ctx, projectId)
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
	}
	return client
}

// The purpose of the particle processor worker is to listen for new particles
// across all streams and perform side effects such as
//   - generate transcript if the particle is of type media
//   - send mobile notifications if a client is offline
//   - update the parent stream's `last_child_created_at`
//   - generate vector embedding
//   - synthesize and decide whether ai should generate a particle as a response
func main() {
	ctx := context.Background()

	db.Init()
	defer db.Cleanup()
	processingRepo := particle.NewProcessingRepository(db.Pool())
	billingSvc := billing.NewServiceForWorker(db.Pool())

	storageClient, err := storage.NewClient(ctx)
	if err != nil {
		slog.Error("failed to create GCS client", "error", err)
		os.Exit(1)
	}
	defer storageClient.Close()

	gcsBucket := utils.MustGetEnv("GCS_BUCKET")
	depotSvc := depot.NewService(db.Pool(), storageClient, depot.Config{
		GoogleServiceAccountEmail: utils.MustGetEnv("GOOGLE_SERVICE_ACCOUNT_EMAIL"),
		BucketName:                gcsBucket,
	})

	speechSvc := speech.NewSpeechService(ctx)

	client := createClient(ctx)
	defer client.Close()

	cutoff := time.Now().Add(-5 * time.Minute)
	it := client.CollectionGroup("children").
		Where("created_at", ">", cutoff).
		Snapshots(ctx)

	for {
		snap, err := it.Next()
		if e := status.Code(err); e == codes.DeadlineExceeded || e == codes.Canceled {
			panic(fmt.Errorf("error: %w", err))
		}
		if err != nil {
			slog.Error("error in processing snapshot", "error", err)
			continue
		}

		if snap == nil {
			continue
		}

		for _, change := range snap.Changes {
			if change.Kind != firestore.DocumentAdded {
				continue
			}

			particleID := change.Doc.Ref.ID

			processed, err := processingRepo.IsProcessed(ctx, particleID)
			if err != nil {
				slog.Error("failed to check processing status", "particleID", particleID, "error", err)
				continue
			}
			if processed {
				slog.Debug("skipping already processed particle", "particleID", particleID)
				continue
			}

			slog.Debug("processing particle", "particleID", particleID, "data", change.Doc.Data())

			// --- Perform side effects ---

			updateParentLastChildCreatedAt(ctx, change.Doc)
			transcribeMediaParticle(ctx, depotSvc, speechSvc, change.Doc)
			transcodeMediaParticle(ctx, depotSvc, change.Doc)
			recordFreemiumUsage(ctx, billingSvc, change.Doc)

			if err := processingRepo.MarkProcessed(ctx, particleID); err != nil {
				slog.Error("failed to mark particle as processed", "particleID", particleID, "error", err)
			}
		}
	}
}

func transcribeMediaParticle(ctx context.Context, depotSvc depot.Service, speechSvc speech.SpeechService, doc *firestore.DocumentSnapshot) {
	var mediaParticle particle.FirestoreMediaParticle
	err := doc.DataTo(&mediaParticle)
	if err != nil {
		slog.Error("unable to marshal particle data", "error", err)
		return
	}

	particleType, err := particle.ParseParticleType(mediaParticle.Type)
	if err != nil {
		slog.Error("invalid particle type", "error", err)
		return
	}

	if particleType != particle.TypeMedia {
		slog.Info("received a particle of type", "particle type", particleType)
		return
	}

	downloadURL, err := depotSvc.GetDownloadURL(ctx, mediaParticle.Properties.ObjectId)
	if err != nil {
		slog.Error("failed to get download URL", "error", err, "object_id", mediaParticle.Properties.ObjectId)
		return
	}

	result, err := speechSvc.Transcribe(ctx, downloadURL)
	if err != nil {
		slog.Error("failed to transcribe media", "error", err, "particleID", doc.Ref.ID)
		return
	}

	transcript := toFirestoreTranscript(result)

	_, err = doc.Ref.Set(ctx, map[string]interface{}{
		"properties": map[string]interface{}{
			"transcript": transcript,
		},
	}, firestore.MergeAll)
	if err != nil {
		slog.Error("failed to update transcript in firestore", "error", err, "particleID", doc.Ref.ID)
		return
	}

	slog.Info("transcribed media particle", "particleID", doc.Ref.ID)
}

// transcodeMediaParticle produces an iOS-playable MP4/m4a derivative for media
// particles whose original mime type AVPlayer can't decode (notably the WebM
// the desktop recorder emits today). Skips when the source is already in an
// iOS-playable family or when a transcoded variant has already been written.
//
// ffmpeg reads directly from the GCS signed URL and writes to a local temp
// file — `+faststart` requires seekable output, so a stdout pipe wouldn't work.
// The temp file is then streamed to GCS via depotSvc.CreateFromReader (no
// presigned-PUT round-trip — the worker has direct SDK access).
func transcodeMediaParticle(ctx context.Context, depotSvc depot.Service, doc *firestore.DocumentSnapshot) {
	var mediaParticle particle.FirestoreMediaParticle
	if err := doc.DataTo(&mediaParticle); err != nil {
		slog.Error("transcode: unable to marshal particle data", "error", err)
		return
	}

	particleType, err := particle.ParseParticleType(mediaParticle.Type)
	if err != nil {
		slog.Error("transcode: invalid particle type", "error", err)
		return
	}
	if particleType != particle.TypeMedia {
    slog.Info("transcode: particle is not of type media")
		return
	}

	// Already transcoded — re-delivery within the 5-min Firestore window.
	if mediaParticle.Properties.TranscodedObjectId != "" {
		return
	}

	// Source is already iOS-playable; nothing to do.
	if isIOSPlayableMime(mediaParticle.Properties.MimeType) {
    slog.Info("transcode: skipping because already playable on ios")
		return
	}

	sourceURL, err := depotSvc.GetDownloadURL(ctx, mediaParticle.Properties.ObjectId)
	if err != nil {
		slog.Error("transcode: failed to get download URL", "error", err, "object_id", mediaParticle.Properties.ObjectId)
		return
	}

	isAudio := strings.HasPrefix(mediaParticle.Properties.MimeType, "audio/")
	var outputExt, outputMime string
	if isAudio {
		outputExt = ".m4a"
		outputMime = "audio/mp4"
	} else {
		outputExt = ".mp4"
		outputMime = "video/mp4"
	}

	tmp, err := os.CreateTemp("", "transcode-*"+outputExt)
	if err != nil {
		slog.Error("transcode: failed to create temp file", "error", err)
		return
	}
	tmpPath := tmp.Name()
	tmp.Close()
	defer os.Remove(tmpPath)

	transcodeCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()

	var args []string
	if isAudio {
		args = []string{
			"-y", "-i", sourceURL,
			"-vn",
			"-c:a", "aac", "-b:a", "128k",
			"-movflags", "+faststart",
			tmpPath,
		}
	} else {
		// Cap encoder parallelism and lookahead to keep memory bounded — screen
		// recordings come in at native display resolution (often 1440p–4K) and
		// libx264's per-thread lookahead/reference buffers blow past the worker's
		// memory limit otherwise. Output is also downscaled to 1080p max, which
		// mobile playback won't notice; the original WebM stays in GCS untouched.
		args = []string{
			"-y", "-i", sourceURL,
			"-vf", "scale='min(1920,iw)':-2:flags=lanczos",
			"-c:v", "libx264", "-preset", "veryfast", "-crf", "23",
			"-pix_fmt", "yuv420p", "-profile:v", "baseline", "-level", "3.1",
			"-x264-params", "rc-lookahead=20:ref=2",
			"-threads", "2", "-filter_threads", "2",
			"-c:a", "aac", "-b:a", "128k",
			"-movflags", "+faststart",
			tmpPath,
		}
	}

	cmd := exec.CommandContext(transcodeCtx, "ffmpeg", args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		slog.Error("transcode: ffmpeg failed", "error", err, "stderr", stderr.String(), "particleID", doc.Ref.ID)
		return
	}

	networkID, err := networkIDFromParticlePath(doc.Ref.Path)
	if err != nil {
		slog.Error("transcode: failed to derive network id", "error", err, "path", doc.Ref.Path)
		return
	}

	f, err := os.Open(tmpPath)
	if err != nil {
		slog.Error("transcode: failed to open transcoded file", "error", err, "path", tmpPath)
		return
	}
	defer f.Close()

	newObj, err := depotSvc.CreateFromReader(ctx, depot.CreateFromReaderInput{
		Prefix:      networkID,
		Name:        "transcoded" + outputExt,
		ContentType: outputMime,
	}, f)
	if err != nil {
		slog.Error("transcode: failed to upload transcoded object", "error", err, "particleID", doc.Ref.ID)
		return
	}

	_, err = doc.Ref.Set(ctx, map[string]interface{}{
		"properties": map[string]interface{}{
			"transcoded_object_id": newObj.ID,
			"transcoded_mime_type": outputMime,
		},
	}, firestore.MergeAll)
	if err != nil {
		slog.Error("transcode: failed to update particle in firestore", "error", err, "particleID", doc.Ref.ID)
		return
	}

	slog.Info("transcoded media particle", "particleID", doc.Ref.ID, "transcoded_object_id", newObj.ID, "mime", outputMime)
}

func isIOSPlayableMime(mime string) bool {
	switch mime {
	case "video/mp4", "video/quicktime", "audio/mp4", "audio/aac", "audio/x-m4a", "audio/mpeg":
		return true
	}
	return false
}

func toFirestoreTranscript(result *speech.TranscriptResult) particle.FirestoreTranscript {
	words := make([]particle.FirestoreTranscriptWord, len(result.Words))
	for i, w := range result.Words {
		words[i] = particle.FirestoreTranscriptWord{
			Word:  w.Word,
			Start: w.Start,
			End:   w.End,
		}
	}

	paragraphs := make([]particle.FirestoreTranscriptParagraph, len(result.Paragraphs))
	for i, p := range result.Paragraphs {
		sentences := make([]particle.FirestoreTranscriptSentence, len(p.Sentences))
		for j, s := range p.Sentences {
			sentences[j] = particle.FirestoreTranscriptSentence{
				Text:  s.Text,
				Start: s.Start,
				End:   s.End,
			}
		}
		paragraphs[i] = particle.FirestoreTranscriptParagraph{
			Sentences: sentences,
			Start:     p.Start,
			End:       p.End,
		}
	}

	return particle.FirestoreTranscript{
		Transcript: result.Transcript,
		Words:      words,
		Paragraphs: paragraphs,
	}
}

// recordFreemiumUsage bumps the network's daily message counter for non-container
// particles. Idempotent via the surrounding processed_particles guard: the worker
// only reaches this path on first-seen particles, so a crash/restart won't
// double-count.
func recordFreemiumUsage(ctx context.Context, billingSvc billing.Service, doc *firestore.DocumentSnapshot) {
	rawType, err := doc.DataAt("type")
	if err != nil {
		slog.Error("failed to read particle type", "error", err, "particleID", doc.Ref.ID)
		return
	}
	typeStr, ok := rawType.(string)
	if !ok {
		slog.Error("particle type is not a string", "particleID", doc.Ref.ID, "type", rawType)
		return
	}
	particleType, err := particle.ParseParticleType(typeStr)
	if err != nil {
		slog.Error("invalid particle type", "error", err, "particleID", doc.Ref.ID)
		return
	}
	// Containers (stream/folder) don't count as "messages" for the daily cap.
	if particleType == particle.TypeStream || particleType == particle.TypeFolder {
		return
	}

	networkID, err := networkIDFromParticlePath(doc.Ref.Path)
	if err != nil {
		slog.Error("failed to derive network id", "error", err, "path", doc.Ref.Path)
		return
	}

	if err := billingSvc.IncrementDailyUsage(ctx, networkID, doc.CreateTime); err != nil {
		slog.Error("failed to increment daily usage", "error", err, "networkID", networkID, "particleID", doc.Ref.ID)
	}
}

// networkIDFromParticlePath extracts the network id from a Firestore particle
// document path. Particles live at `networks/{network_id}/children/.../children/{id}`
// at arbitrary nesting depth, so the network id is always the second segment
// of the full doc path (which itself is rooted under the Firestore db path:
// `projects/.../documents/networks/{network_id}/...`).
func networkIDFromParticlePath(path string) (string, error) {
	// doc.Ref.Path is the full resource path; find the "networks" collection
	// and return the next segment.
	segments := strings.Split(path, "/")
	for i, seg := range segments {
		if seg == "networks" && i+1 < len(segments) {
			return segments[i+1], nil
		}
	}
	return "", fmt.Errorf("no networks segment in path: %s", path)
}

// updateParentLastChildCreatedAt updates the parent stream's last_child_created_at
// to the child's actual created_at timestamp, so it stays directly comparable with
// playback markers (which also store child created_at values).
func updateParentLastChildCreatedAt(ctx context.Context, doc *firestore.DocumentSnapshot) {
	parentChildrenCollectionRef := doc.Ref.Parent
	if parentChildrenCollectionRef == nil {
		return
	}

	parentParticleDocRef := parentChildrenCollectionRef.Parent
	if parentParticleDocRef == nil {
		slog.Error("particle has no parent document", "particleID", doc.Ref.ID)
		return
	}

	parentParticleDoc, err := parentParticleDocRef.Get(ctx)
	if err != nil {
		slog.Error("failed to get parent particle", "error", err)
		return
	}

	slog.Info("parent particle is", "parent particle id", parentParticleDoc.Ref.ID)

	var streamParticle particle.FirestoreStreamParticle
	if err := parentParticleDoc.DataTo(&streamParticle); err != nil {
		slog.Error("failed to parse stream particle", "error", err)
		return
	}

	particleType, err := particle.ParseParticleType(streamParticle.Type)
	if err != nil {
		slog.Error("invalid particle type", "error", err)
		return
	}

	if particleType != particle.TypeStream {
		return
	}

	// Read the child's created_at — this is the same value that playback markers store
	childCreatedAt, err := doc.DataAt("created_at")
	if err != nil {
		slog.Error("failed to read child created_at", "error", err, "particleID", doc.Ref.ID)
		return
	}

	slog.Info("going to update the last_child_created_at for parent particle")
	_, err = parentParticleDocRef.Update(ctx, []firestore.Update{
		{
			Path:  "last_child_created_at",
			Value: childCreatedAt,
		},
	})
	if err != nil {
		slog.Error("unable to update parent particle `last_child_created_at`")
	}
}