core/server/handlers/completion_chain.go

package handlers

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"chat-switchboard/events"
	"chat-switchboard/models"
	"chat-switchboard/providers"

	capspkg "chat-switchboard/capabilities"
)

// maxChainDepth limits AI-to-AI chaining to prevent runaway loops.
const maxChainDepth = 5

// chainIfMentioned checks if an assistant response @mentions another
// persona and triggers a follow-up completion. Uses resolveMention()
// directly — no roster or participant list needed. Works in any chat.
//
// Runs asynchronously (goroutine) after the SSE stream closes.
// The follow-up response is delivered to the client via WebSocket events.
func (h *CompletionHandler) chainIfMentioned(
	channelID, userID, currentPersonaID, responseContent string,
	depth int,
) {
	if depth >= maxChainDepth {
		return
	}

	defer func() {
		if r := recover(); r != nil {
			log.Printf("[chain] PANIC in chain (depth=%d): %v", depth, r)
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
	defer cancel()

	// Extract the @token for logging
	token := extractFirstMention(responseContent)
	if token == "" {
		log.Printf("[chain] No @mention found in response (len=%d, persona=%s)", len(responseContent), currentPersonaID[:min(8, len(currentPersonaID))])
		return
	}
	log.Printf("[chain] Found @%s in response from persona %s (depth=%d)", token, currentPersonaID[:min(8, len(currentPersonaID))], depth)

	// Use the same resolveMention() that handles user @mentions
	mentionModel, mentionConfig, mentionPersona, _ := h.resolveMention(ctx, userID, responseContent)
	if mentionModel == "" {
		log.Printf("[chain] @%s did not resolve to any model", token)
		return
	}
	if mentionPersona == nil {
		log.Printf("[chain] @%s resolved to raw model %s (no chain for raw models)", token, mentionModel)
		return // no persona @mention found, or it's a raw model (no chain for those)
	}

	// Don't chain to self
	if mentionPersona.ID == currentPersonaID {
		log.Printf("[chain] @%s is self-mention, skipping", token)
		return
	}

	log.Printf("[chain] %s @mentioned %s (depth=%d)", currentPersonaID[:min(8, len(currentPersonaID))], mentionPersona.Name, depth+1)

	// Emit typing indicator
	if h.hub != nil {
		h.emitTyping(userID, channelID, mentionPersona.ID, mentionPersona.Name, true)
		defer h.emitTyping(userID, channelID, mentionPersona.ID, mentionPersona.Name, false)
	}

	// Brief pause for UX
	time.Sleep(500 * time.Millisecond)

	// Resolve provider config
	configID := mentionConfig
	chainReq := completionRequest{
		ChannelID:        channelID,
		Model:            mentionModel,
		ProviderConfigID: configID,
	}
	providerCfg, providerID, model, resolvedConfigID, _, err := h.resolveConfig(userID, channelID, chainReq)
	if err != nil {
		log.Printf("[chain] Failed to resolve config for %s: %v", mentionPersona.Name, err)
		return
	}
	configID = resolvedConfigID

	provider, err := providers.Get(providerID)
	if err != nil {
		log.Printf("[chain] Provider %s unavailable: %v", providerID, err)
		return
	}

	// Load conversation with persona's system prompt
	messages, err := h.loadConversation(channelID, userID, mentionPersona.SystemPrompt, mentionPersona.ID, model)
	if err != nil {
		log.Printf("[chain] Failed to load conversation: %v", err)
		return
	}

	// Context boundary
	boundary := fmt.Sprintf(
		"[You are now %s. Previous assistant messages may be from different AI personas — ignore their style and mannerisms. Respond only as yourself per your system prompt.]",
		mentionPersona.Name,
	)
	messages = append(messages, providers.Message{Role: "system", Content: boundary})

	caps := capspkg.ResolveIntrinsic(model, nil, nil)

	provReq := providers.CompletionRequest{
		Model:     model,
		Messages:  messages,
		MaxTokens: capspkg.ResolveMaxOutput(model, caps),
	}
	if mentionPersona.Temperature != nil {
		provReq.Temperature = mentionPersona.Temperature
	}
	if mentionPersona.ThinkingBudget != nil && *mentionPersona.ThinkingBudget > 0 {
		if providerCfg.Settings == nil {
			providerCfg.Settings = make(map[string]interface{})
		}
		providerCfg.Settings["extended_thinking"] = true
		providerCfg.Settings["thinking_budget"] = *mentionPersona.ThinkingBudget
	}

	// Apply provider-specific hooks
	if hooks := providers.GetHooks(providerID); hooks != nil {
		hooks.PreRequest(providerCfg, &provReq)
	}

	// Synchronous completion — result delivered via WebSocket
	callStart := time.Now()
	resp, err := provider.ChatCompletion(ctx, providerCfg, provReq)
	latencyMs := int(time.Since(callStart).Milliseconds())
	if err != nil {
		log.Printf("[chain] Completion failed for %s: %v", mentionPersona.Name, err)
		if h.health != nil {
			h.health.RecordError(configID, latencyMs, err.Error())
		}
		return
	}
	if h.health != nil {
		h.health.RecordSuccess(configID, latencyMs)
	}

	if resp.Content == "" {
		return
	}

	// Persist
	msgID, err := h.persistMessage(channelID, userID, "assistant", resp.Content, model,
		resp.InputTokens, resp.OutputTokens, nil, nil, configID, mentionPersona.ID)
	if err != nil {
		log.Printf("[chain] Failed to persist chained message: %v", err)
		return
	}

	// Deliver via WebSocket to all channel participants
	if h.hub != nil {
		broadcastAssistantMessage(context.Background(), h.hub, channelID, msgID, resp.Content, model, "", mentionPersona.ID)
	}

	// Log usage
	if h.stores.Usage != nil {
		entry := &models.UsageEntry{
			ChannelID:        &channelID,
			UserID:           userID,
			ProviderConfigID: &configID,
			ProviderScope:    "global",
			ModelID:          model,
			PromptTokens:     resp.InputTokens,
			CompletionTokens: resp.OutputTokens,
		}
		_ = h.stores.Usage.Log(ctx, entry)
	}

	log.Printf("[chain] ✅ %s responded (depth=%d, %d tokens) in channel %s",
		mentionPersona.Name, depth+1, resp.InputTokens+resp.OutputTokens, channelID[:min(8, len(channelID))])

	// Recursive: check if this response @mentions yet another persona
	h.chainIfMentioned(channelID, userID, mentionPersona.ID, resp.Content, depth+1)
}

// chainToPersona triggers a completion for a specific persona by ID.
// Used by @all fan-out. Does NOT recurse (depth-1 only).
func (h *CompletionHandler) chainToPersona(channelID, userID, personaID, triggerContent string, depth int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[chain-all] PANIC: %v", r)
		}
	}()

	persona := ResolvePersona(h.stores, personaID, userID)
	if persona == nil {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
	defer cancel()

	if h.hub != nil {
		h.emitTyping(userID, channelID, persona.ID, persona.Name, true)
		defer h.emitTyping(userID, channelID, persona.ID, persona.Name, false)
	}

	configID := ""
	if persona.ProviderConfigID != nil {
		configID = *persona.ProviderConfigID
	}
	chainReq := completionRequest{
		ChannelID:        channelID,
		Model:            persona.BaseModelID,
		ProviderConfigID: configID,
	}
	providerCfg, providerID, model, resolvedConfigID, _, err := h.resolveConfig(userID, channelID, chainReq)
	if err != nil {
		log.Printf("[chain-all] Config failed for %s: %v", persona.Name, err)
		return
	}
	configID = resolvedConfigID

	provider, err := providers.Get(providerID)
	if err != nil {
		return
	}

	messages, err := h.loadConversation(channelID, userID, persona.SystemPrompt, persona.ID, model)
	if err != nil {
		return
	}

	boundary := fmt.Sprintf(
		"[You are now %s. Previous assistant messages may be from different AI personas — ignore their style. Respond only as yourself per your system prompt.]",
		persona.Name,
	)
	messages = append(messages, providers.Message{Role: "system", Content: boundary})

	caps := capspkg.ResolveIntrinsic(model, nil, nil)
	provReq := providers.CompletionRequest{
		Model:     model,
		Messages:  messages,
		MaxTokens: capspkg.ResolveMaxOutput(model, caps),
	}
	if persona.Temperature != nil {
		provReq.Temperature = persona.Temperature
	}
	if persona.ThinkingBudget != nil && *persona.ThinkingBudget > 0 {
		if providerCfg.Settings == nil {
			providerCfg.Settings = make(map[string]interface{})
		}
		providerCfg.Settings["extended_thinking"] = true
		providerCfg.Settings["thinking_budget"] = *persona.ThinkingBudget
	}

	if hooks := providers.GetHooks(providerID); hooks != nil {
		hooks.PreRequest(providerCfg, &provReq)
	}

	callStart := time.Now()
	resp, err := provider.ChatCompletion(ctx, providerCfg, provReq)
	latencyMs := int(time.Since(callStart).Milliseconds())
	if err != nil {
		if h.health != nil {
			h.health.RecordError(configID, latencyMs, err.Error())
		}
		return
	}
	if h.health != nil {
		h.health.RecordSuccess(configID, latencyMs)
	}
	if resp.Content == "" {
		return
	}

	msgID, err := h.persistMessage(channelID, userID, "assistant", resp.Content, model,
		resp.InputTokens, resp.OutputTokens, nil, nil, configID, persona.ID)
	if err != nil {
		return
	}

	if h.hub != nil {
		broadcastAssistantMessage(context.Background(), h.hub, channelID, msgID, resp.Content, model, "", persona.ID)
	}

	if h.stores.Usage != nil {
		entry := &models.UsageEntry{
			ChannelID:        &channelID,
			UserID:           userID,
			ProviderConfigID: &configID,
			ProviderScope:    "global",
			ModelID:          model,
			PromptTokens:     resp.InputTokens,
			CompletionTokens: resp.OutputTokens,
		}
		_ = h.stores.Usage.Log(ctx, entry)
	}

	log.Printf("[chain-all] ✅ %s responded (%d tokens) in channel %s",
		persona.Name, resp.InputTokens+resp.OutputTokens, channelID[:min(8, len(channelID))])
}

// emitTyping sends a typing indicator via WebSocket.
func (h *CompletionHandler) emitTyping(userID, channelID, participantID, displayName string, start bool) {
	eventType := "typing.start"
	if !start {
		eventType = "typing.stop"
	}
	payload, _ := json.Marshal(map[string]string{
		"channel_id":       channelID,
		"participant_id":   participantID,
		"participant_type": "persona",
		"display_name":     displayName,
	})
	h.hub.PublishToUser(userID, events.Event{
		Label:   eventType,
		Payload: payload,
		Ts:      time.Now().UnixMilli(),
	})
}