kandev/backend/internal/orchestrator/watcher/watcher.go

// Package watcher provides event subscription and dispatching for the Orchestrator.
package watcher

import (
	"context"
	"encoding/json"
	"strings"
	"sync"

	"go.uber.org/zap"

	"github.com/kandev/kandev/internal/common/logger"
	"github.com/kandev/kandev/internal/events"
	"github.com/kandev/kandev/internal/events/bus"
	v1 "github.com/kandev/kandev/pkg/api/v1"
	"github.com/kandev/kandev/pkg/acp/protocol"
)

// TaskEventData contains data from task events
type TaskEventData struct {
	TaskID   string        `json:"task_id"`
	Task     *v1.Task      `json:"task,omitempty"`
	OldState *v1.TaskState `json:"old_state,omitempty"`
	NewState *v1.TaskState `json:"new_state,omitempty"`
}

// AgentEventData contains data from agent events
type AgentEventData struct {
	TaskID          string `json:"task_id"`
	AgentInstanceID string `json:"agent_instance_id"`
	AgentType       string `json:"agent_type"`
	ExitCode        *int   `json:"exit_code,omitempty"`
	ErrorMessage    string `json:"error_message,omitempty"`
}

// EventHandlers contains callbacks for different event types
type EventHandlers struct {
	// Task events
	OnTaskCreated      func(ctx context.Context, data TaskEventData)
	OnTaskUpdated      func(ctx context.Context, data TaskEventData)
	OnTaskStateChanged func(ctx context.Context, data TaskEventData)
	OnTaskDeleted      func(ctx context.Context, data TaskEventData)

	// Agent events
	OnAgentStarted   func(ctx context.Context, data AgentEventData)
	OnAgentCompleted func(ctx context.Context, data AgentEventData)
	OnAgentFailed    func(ctx context.Context, data AgentEventData)
	OnAgentStopped   func(ctx context.Context, data AgentEventData)

	// ACP messages
	OnACPMessage func(ctx context.Context, taskID string, msg *protocol.Message)
}

// Watcher subscribes to events and dispatches to handlers
type Watcher struct {
	eventBus bus.EventBus
	handlers EventHandlers
	logger   *logger.Logger

	subscriptions []bus.Subscription
	mu            sync.Mutex
	running       bool
}

// queueName is the queue group for load balancing across orchestrator instances
const queueName = "orchestrator"

// NewWatcher creates a new event watcher
func NewWatcher(eventBus bus.EventBus, handlers EventHandlers, log *logger.Logger) *Watcher {
	return &Watcher{
		eventBus:      eventBus,
		handlers:      handlers,
		logger:        log.WithFields(zap.String("component", "watcher")),
		subscriptions: make([]bus.Subscription, 0),
	}
}

// Start begins watching for events
func (w *Watcher) Start(ctx context.Context) error {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.running {
		return nil
	}

	w.logger.Info("Starting event watcher")

	// Subscribe to task events
	if err := w.subscribeToTaskEvents(); err != nil {
		return err
	}

	// Subscribe to agent events
	if err := w.subscribeToAgentEvents(); err != nil {
		w.unsubscribeAll()
		return err
	}

	// Subscribe to ACP messages
	if err := w.subscribeToACPMessages(); err != nil {
		w.unsubscribeAll()
		return err
	}

	w.running = true
	w.logger.Info("Event watcher started", zap.Int("subscriptions", len(w.subscriptions)))
	return nil
}

// Stop stops watching for events
func (w *Watcher) Stop() error {
	w.mu.Lock()
	defer w.mu.Unlock()

	if !w.running {
		return nil
	}

	w.logger.Info("Stopping event watcher")
	w.unsubscribeAll()
	w.running = false
	w.logger.Info("Event watcher stopped")
	return nil
}

// IsRunning returns true if the watcher is active
func (w *Watcher) IsRunning() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.running
}

// unsubscribeAll removes all subscriptions (must be called with lock held)
func (w *Watcher) unsubscribeAll() {
	for _, sub := range w.subscriptions {
		if err := sub.Unsubscribe(); err != nil {
			w.logger.Error("Failed to unsubscribe", zap.Error(err))
		}
	}
	w.subscriptions = make([]bus.Subscription, 0)
}

// subscribeToTaskEvents subscribes to all task events
func (w *Watcher) subscribeToTaskEvents() error {
	taskEvents := []struct {
		subject string
		handler func(ctx context.Context, data TaskEventData)
	}{
		{events.TaskCreated, w.handlers.OnTaskCreated},
		{events.TaskUpdated, w.handlers.OnTaskUpdated},
		{events.TaskStateChanged, w.handlers.OnTaskStateChanged},
		{events.TaskDeleted, w.handlers.OnTaskDeleted},
	}

	for _, te := range taskEvents {
		if te.handler == nil {
			continue
		}
		handler := te.handler // capture for closure
		sub, err := w.eventBus.QueueSubscribe(te.subject, queueName, w.createTaskEventHandler(handler))
		if err != nil {
			w.logger.Error("Failed to subscribe to task event",
				zap.String("subject", te.subject),
				zap.Error(err))
			return err
		}
		w.subscriptions = append(w.subscriptions, sub)
	}
	return nil
}

// subscribeToAgentEvents subscribes to all agent events
func (w *Watcher) subscribeToAgentEvents() error {
	agentEvents := []struct {
		subject string
		handler func(ctx context.Context, data AgentEventData)
	}{
		{events.AgentStarted, w.handlers.OnAgentStarted},
		{events.AgentCompleted, w.handlers.OnAgentCompleted},
		{events.AgentFailed, w.handlers.OnAgentFailed},
		{events.AgentStopped, w.handlers.OnAgentStopped},
	}

	for _, ae := range agentEvents {
		if ae.handler == nil {
			continue
		}
		handler := ae.handler // capture for closure
		sub, err := w.eventBus.QueueSubscribe(ae.subject, queueName, w.createAgentEventHandler(handler))
		if err != nil {
			w.logger.Error("Failed to subscribe to agent event",
				zap.String("subject", ae.subject),
				zap.Error(err))
			return err
		}
		w.subscriptions = append(w.subscriptions, sub)
	}
	return nil
}

// subscribeToACPMessages subscribes to ACP messages using wildcard
func (w *Watcher) subscribeToACPMessages() error {
	if w.handlers.OnACPMessage == nil {
		return nil
	}

	// Use wildcard to subscribe to all ACP messages (acp.message.*)
	subject := events.BuildACPWildcardSubject()

	// Use regular subscription for ACP messages (each instance needs all messages for WebSocket streaming)
	sub, err := w.eventBus.Subscribe(subject, w.createACPMessageHandler())
	if err != nil {
		w.logger.Error("Failed to subscribe to ACP messages",
			zap.String("subject", subject),
			zap.Error(err))
		return err
	}
	w.subscriptions = append(w.subscriptions, sub)
	return nil
}

// createTaskEventHandler creates a bus.EventHandler for task events
func (w *Watcher) createTaskEventHandler(handler func(ctx context.Context, data TaskEventData)) bus.EventHandler {
	return func(ctx context.Context, event *bus.Event) error {
		var data TaskEventData
		if err := w.parseEventData(event.Data, &data); err != nil {
			w.logger.Error("Failed to parse task event data",
				zap.String("event_type", event.Type),
				zap.String("event_id", event.ID),
				zap.Error(err))
			return nil // Don't return error to continue processing other events
		}

		w.logger.Debug("Handling task event",
			zap.String("event_type", event.Type),
			zap.String("task_id", data.TaskID))

		handler(ctx, data)
		return nil
	}
}

// createAgentEventHandler creates a bus.EventHandler for agent events
func (w *Watcher) createAgentEventHandler(handler func(ctx context.Context, data AgentEventData)) bus.EventHandler {
	return func(ctx context.Context, event *bus.Event) error {
		var data AgentEventData
		if err := w.parseEventData(event.Data, &data); err != nil {
			w.logger.Error("Failed to parse agent event data",
				zap.String("event_type", event.Type),
				zap.String("event_id", event.ID),
				zap.Error(err))
			return nil // Don't return error to continue processing other events
		}

		w.logger.Debug("Handling agent event",
			zap.String("event_type", event.Type),
			zap.String("task_id", data.TaskID),
			zap.String("agent_instance_id", data.AgentInstanceID))

		handler(ctx, data)
		return nil
	}
}

// createACPMessageHandler creates a bus.EventHandler for ACP messages
func (w *Watcher) createACPMessageHandler() bus.EventHandler {
	return func(ctx context.Context, event *bus.Event) error {
		// Parse the ACP message from event data
		var msg protocol.Message
		if err := w.parseEventData(event.Data, &msg); err != nil {
			w.logger.Error("Failed to parse ACP message",
				zap.String("event_id", event.ID),
				zap.Error(err))
			return nil // Don't return error to continue processing other events
		}

		// Extract task ID from subject (acp.message.{task_id})
		taskID := msg.TaskID
		if taskID == "" {
			// Try to extract from event type if not in message
			parts := strings.Split(event.Type, ".")
			if len(parts) >= 3 {
				taskID = parts[2]
			}
		}

		if taskID == "" {
			w.logger.Warn("ACP message missing task_id",
				zap.String("event_id", event.ID))
			return nil
		}

		w.logger.Debug("Handling ACP message",
			zap.String("task_id", taskID),
			zap.String("message_type", string(msg.Type)))

		w.handlers.OnACPMessage(ctx, taskID, &msg)
		return nil
	}
}

// parseEventData converts event data map to a typed struct
func (w *Watcher) parseEventData(data map[string]interface{}, target interface{}) error {
	// Marshal to JSON and unmarshal to target type
	jsonData, err := json.Marshal(data)
	if err != nil {
		return err
	}
	return json.Unmarshal(jsonData, target)
}