geo/server/internal/shared/stream/generation.go

package stream

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
)

type GenerationEvent struct {
	Type      string    `json:"type"`
	ArticleID int64     `json:"article_id"`
	TaskID    int64     `json:"task_id"`
	Title     string    `json:"title,omitempty"`
	Status    string    `json:"status,omitempty"`
	Delta     string    `json:"delta,omitempty"`
	Content   string    `json:"content,omitempty"`
	Error     string    `json:"error,omitempty"`
	Done      bool      `json:"done,omitempty"`
	UpdatedAt time.Time `json:"updated_at"`
}

type generationEventEnvelope struct {
	Source string          `json:"source"`
	Event  GenerationEvent `json:"event"`
}

type generationStream struct {
	snapshot       GenerationEvent
	subscribers    map[int]chan GenerationEvent
	nextSubscriber int
	completed      bool
}

type GenerationHub struct {
	mu           sync.Mutex
	streams      map[int64]*generationStream
	rabbitMQ     *rabbitmq.Client
	instanceID   string
	consumerName string
}

func NewGenerationHub(rabbitMQClient *rabbitmq.Client) *GenerationHub {
	instanceID := fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano())
	return &GenerationHub{
		streams:      make(map[int64]*generationStream),
		rabbitMQ:     rabbitMQClient,
		instanceID:   instanceID,
		consumerName: "generation-stream-" + instanceID,
	}
}

func (h *GenerationHub) Run(ctx context.Context) {
	if h == nil || h.rabbitMQ == nil {
		return
	}
	go h.run(ctx)
}

func (h *GenerationHub) Seed(event GenerationEvent) {
	h.mu.Lock()
	defer h.mu.Unlock()

	state := h.ensureLocked(event.ArticleID)
	applySnapshotLocked(state, event)
	if state.completed && len(state.subscribers) == 0 {
		delete(h.streams, event.ArticleID)
	}
}

func (h *GenerationHub) StartEvent(articleID, taskID int64, title string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.completed = false
		state.snapshot = GenerationEvent{
			Type:      "snapshot",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "generating",
			UpdatedAt: now,
		}
		return GenerationEvent{
			Type:      "status",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "generating",
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Start(articleID, taskID int64, title string) {
	h.StartEvent(articleID, taskID, title)
}

func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.snapshot.Type = "snapshot"
		state.snapshot.ArticleID = articleID
		state.snapshot.TaskID = taskID
		state.snapshot.Title = title
		state.snapshot.Status = "generating"
		state.snapshot.Content += delta
		state.snapshot.UpdatedAt = now
		state.completed = false

		return GenerationEvent{
			Type:      "delta",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "generating",
			Delta:     delta,
			Content:   state.snapshot.Content,
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Complete(articleID, taskID int64, title, content string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.completed = true
		state.snapshot = GenerationEvent{
			Type:      "snapshot",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "completed",
			Content:   content,
			Done:      true,
			UpdatedAt: now,
		}

		return GenerationEvent{
			Type:      "completed",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "completed",
			Content:   content,
			Done:      true,
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Fail(articleID, taskID int64, title, errMsg string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.completed = true
		state.snapshot.Type = "snapshot"
		state.snapshot.ArticleID = articleID
		state.snapshot.TaskID = taskID
		state.snapshot.Title = title
		state.snapshot.Status = "failed"
		state.snapshot.Error = errMsg
		state.snapshot.Done = true
		state.snapshot.UpdatedAt = now

		return GenerationEvent{
			Type:      "error",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "failed",
			Error:     errMsg,
			Done:      true,
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Subscribe(articleID int64) (<-chan GenerationEvent, *GenerationEvent, func()) {
	h.mu.Lock()
	state := h.streams[articleID]
	if state == nil {
		h.mu.Unlock()
		return nil, nil, func() {}
	}

	id := state.nextSubscriber
	state.nextSubscriber++

	ch := make(chan GenerationEvent, 64)
	state.subscribers[id] = ch
	snapshot := state.snapshot
	h.mu.Unlock()

	cancel := func() {
		h.mu.Lock()
		defer h.mu.Unlock()

		current := h.streams[articleID]
		if current == nil {
			close(ch)
			return
		}

		delete(current.subscribers, id)
		shouldDelete := current.completed && len(current.subscribers) == 0
		if shouldDelete {
			delete(h.streams, articleID)
		}

		close(ch)
	}

	return ch, &snapshot, cancel
}

func (h *GenerationHub) Snapshot(articleID int64) (*GenerationEvent, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()

	state := h.streams[articleID]
	if state == nil {
		return nil, false
	}

	snapshot := state.snapshot
	return &snapshot, true
}

func (h *GenerationHub) run(ctx context.Context) {
	for {
		deliveries, ch, err := h.rabbitMQ.ConsumeGenerationEvents(h.consumerName)
		if err != nil {
			select {
			case <-ctx.Done():
				return
			case <-time.After(2 * time.Second):
				continue
			}
		}

		closed := false
		for !closed {
			select {
			case <-ctx.Done():
				_ = ch.Close()
				return
			case delivery, ok := <-deliveries:
				if !ok {
					closed = true
					continue
				}
				h.handleDelivery(delivery.Body)
				_ = delivery.Ack(false)
			}
		}
		_ = ch.Close()

		select {
		case <-ctx.Done():
			return
		case <-time.After(2 * time.Second):
		}
	}
}

func (h *GenerationHub) handleDelivery(body []byte) {
	var envelope generationEventEnvelope
	if err := json.Unmarshal(body, &envelope); err != nil {
		return
	}
	if envelope.Source == h.instanceID {
		return
	}

	h.mu.Lock()
	state := h.streams[envelope.Event.ArticleID]
	if state == nil {
		h.mu.Unlock()
		return
	}

	applyRemoteEventLocked(state, envelope.Event)
	subscribers := collectSubscribersLocked(state)
	shouldDelete := state.completed && len(state.subscribers) == 0
	if shouldDelete {
		delete(h.streams, envelope.Event.ArticleID)
	}
	h.mu.Unlock()

	for _, ch := range subscribers {
		select {
		case ch <- envelope.Event:
		default:
		}
	}
}

func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) GenerationEvent) {
	h.mu.Lock()
	state := h.ensureLocked(articleID)
	event := mutate(state)
	subscribers := collectSubscribersLocked(state)
	shouldDelete := state.completed && len(state.subscribers) == 0
	if shouldDelete {
		delete(h.streams, articleID)
	}
	h.mu.Unlock()

	for _, ch := range subscribers {
		select {
		case ch <- event:
		default:
		}
	}

	if h.rabbitMQ != nil {
		payload, err := json.Marshal(generationEventEnvelope{
			Source: h.instanceID,
			Event:  event,
		})
		if err == nil {
			_ = h.rabbitMQ.PublishGenerationEvent(context.Background(), payload)
		}
	}
}

func (h *GenerationHub) ensureLocked(articleID int64) *generationStream {
	state := h.streams[articleID]
	if state != nil {
		return state
	}

	state = &generationStream{
		subscribers: make(map[int]chan GenerationEvent),
	}
	h.streams[articleID] = state
	return state
}

func collectSubscribersLocked(state *generationStream) []chan GenerationEvent {
	subscribers := make([]chan GenerationEvent, 0, len(state.subscribers))
	for _, ch := range state.subscribers {
		subscribers = append(subscribers, ch)
	}
	return subscribers
}

func applySnapshotLocked(state *generationStream, event GenerationEvent) {
	state.snapshot = GenerationEvent{
		Type:      "snapshot",
		ArticleID: event.ArticleID,
		TaskID:    event.TaskID,
		Title:     event.Title,
		Status:    event.Status,
		Content:   event.Content,
		Error:     event.Error,
		Done:      event.Done,
		UpdatedAt: event.UpdatedAt,
	}
	state.completed = event.Done
}

func applyRemoteEventLocked(state *generationStream, event GenerationEvent) {
	switch event.Type {
	case "delta":
		state.completed = false
		state.snapshot.Type = "snapshot"
		state.snapshot.ArticleID = event.ArticleID
		state.snapshot.TaskID = event.TaskID
		state.snapshot.Title = event.Title
		state.snapshot.Status = event.Status
		if event.Content != "" {
			state.snapshot.Content = event.Content
		} else {
			state.snapshot.Content += event.Delta
		}
		state.snapshot.UpdatedAt = event.UpdatedAt
	case "completed":
		applySnapshotLocked(state, GenerationEvent{
			Type:      "snapshot",
			ArticleID: event.ArticleID,
			TaskID:    event.TaskID,
			Title:     event.Title,
			Status:    "completed",
			Content:   event.Content,
			Done:      true,
			UpdatedAt: event.UpdatedAt,
		})
	case "error":
		applySnapshotLocked(state, GenerationEvent{
			Type:      "snapshot",
			ArticleID: event.ArticleID,
			TaskID:    event.TaskID,
			Title:     event.Title,
			Status:    "failed",
			Error:     event.Error,
			Done:      true,
			UpdatedAt: event.UpdatedAt,
		})
	default:
		applySnapshotLocked(state, GenerationEvent{
			Type:      "snapshot",
			ArticleID: event.ArticleID,
			TaskID:    event.TaskID,
			Title:     event.Title,
			Status:    event.Status,
			Content:   event.Content,
			Error:     event.Error,
			Done:      event.Done,
			UpdatedAt: event.UpdatedAt,
		})
	}
}