core/server/handlers/completion.go

package handlers

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"

	"git.gobha.me/xcaliber/chat-switchboard/database"
	"git.gobha.me/xcaliber/chat-switchboard/providers"
	"git.gobha.me/xcaliber/chat-switchboard/tools"
)

// ── Request Types ───────────────────────────

type completionRequest struct {
	ChannelID   string   `json:"channel_id"` // preferred; validated manually below
	ChatID      string   `json:"chat_id"`    // deprecated alias — maps to channel_id
	Content     string   `json:"content" binding:"required"`
	Model       string   `json:"model,omitempty"`
	PresetID    string   `json:"preset_id,omitempty"` // if set, unwraps preset → base model + config
	APIConfigID string   `json:"api_config_id,omitempty"`
	MaxTokens   int      `json:"max_tokens,omitempty"`
	Temperature *float64 `json:"temperature,omitempty"`
	TopP        *float64 `json:"top_p,omitempty"`
	Stream      *bool    `json:"stream,omitempty"`
}

// CompletionHandler proxies LLM requests through the backend.
type CompletionHandler struct{}

// NewCompletionHandler creates a new handler.
func NewCompletionHandler() *CompletionHandler {
	return &CompletionHandler{}
}

// ── Chat Completion ─────────────────────────
// POST /api/v1/chat/completions
//
// Flow:
//   1. Validate request, verify chat ownership
//   2. Resolve api_config (from request, chat, or user default)
//   3. Load conversation history from messages
//   4. Persist user message
//   5. Attach tool definitions if model supports tools
//   6. Call provider (stream or non-stream)
//   7. Tool loop: if LLM requests tool calls, execute them and re-call
//   8. Stream SSE to client / return JSON
//   9. Persist assistant message with token counts

func (h *CompletionHandler) Complete(c *gin.Context) {
	var req completionRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// Support chat_id as alias during frontend transition
	channelID := req.ChannelID
	if channelID == "" {
		channelID = req.ChatID
	}
	if channelID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "channel_id is required"})
		return
	}

	userID := getUserID(c)

	// Verify channel ownership
	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	// ── Preset unwrap: preset overrides defaults, explicit request fields win ──
	var presetSystemPrompt string
	if req.PresetID != "" {
		preset := ResolvePreset(req.PresetID, userID)
		if preset == nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "preset not found or not accessible"})
			return
		}
		// Preset provides defaults; explicit request fields take priority
		if req.Model == "" {
			req.Model = preset.BaseModelID
		}
		if req.APIConfigID == "" && preset.APIConfigID != nil {
			req.APIConfigID = *preset.APIConfigID
		}
		if req.Temperature == nil && preset.Temperature != nil {
			req.Temperature = preset.Temperature
		}
		if req.MaxTokens == 0 && preset.MaxTokens != nil {
			req.MaxTokens = *preset.MaxTokens
		}
		if preset.SystemPrompt != "" {
			presetSystemPrompt = preset.SystemPrompt
		}
	}

	// Resolve provider config
	providerCfg, providerID, model, configID, err := h.resolveConfig(userID, channelID, req)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// ── Team policy: require_private_providers ──
	if err := enforcePrivateProviderPolicy(userID, configID); err != nil {
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}

	provider, err := providers.Get(providerID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// Load conversation history
	messages, err := h.loadConversation(channelID, userID, presetSystemPrompt)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load conversation"})
		return
	}

	// Add the new user message
	messages = append(messages, providers.Message{
		Role:    "user",
		Content: req.Content,
	})

	// Persist user message
	if _, err := h.persistMessage(channelID, userID, "user", req.Content, "", 0, 0, nil); err != nil {
		log.Printf("Failed to persist user message: %v", err)
	}

	// Build provider request
	provReq := providers.CompletionRequest{
		Model:    model,
		Messages: messages,
	}

	// Resolve capabilities for this model — auto-set defaults
	caps := h.getModelCapabilities(model, configID)

	if req.MaxTokens > 0 {
		provReq.MaxTokens = req.MaxTokens
	} else {
		// ResolveMaxOutput checks: caps → known models → context/8 → 4096
		provReq.MaxTokens = providers.ResolveMaxOutput(model, caps)
	}

	if req.Temperature != nil {
		provReq.Temperature = req.Temperature
	}
	if req.TopP != nil {
		provReq.TopP = req.TopP
	}

	// Attach tool definitions if model supports tool calling and tools are available
	if caps.ToolCalling && tools.HasTools() {
		provReq.Tools = h.buildToolDefs()
	}

	// Determine streaming
	stream := true // default
	if req.Stream != nil {
		stream = *req.Stream
	}

	if stream {
		h.streamCompletion(c, provider, providerCfg, provReq, channelID, userID, model)
	} else {
		h.syncCompletion(c, provider, providerCfg, provReq, channelID, userID, model)
	}
}

// buildToolDefs converts registered tools to provider-format tool definitions.
func (h *CompletionHandler) buildToolDefs() []providers.ToolDef {
	allTools := tools.AllDefinitions()
	defs := make([]providers.ToolDef, len(allTools))
	for i, t := range allTools {
		defs[i] = providers.ToolDef{
			Type: "function",
			Function: providers.FunctionDef{
				Name:        t.Name,
				Description: t.Description,
				Parameters:  t.Parameters,
			},
		}
	}
	return defs
}

// ── Streaming Completion (SSE) with Tool Loop ──

const maxToolIterations = 10 // safety limit on tool call rounds

func (h *CompletionHandler) streamCompletion(
	c *gin.Context,
	provider providers.Provider,
	cfg providers.ProviderConfig,
	req providers.CompletionRequest,
	channelID, userID, model string,
) {
	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no") // Disable nginx buffering

	c.Status(http.StatusOK)
	flusher, _ := c.Writer.(http.Flusher)

	flush := func() {
		if flusher != nil {
			flusher.Flush()
		}
	}

	// sendSSE sends a standard OpenAI-compatible SSE data line.
	sendSSE := func(data string) {
		fmt.Fprintf(c.Writer, "data: %s\n\n", data)
		flush()
	}

	// sendEvent sends a named SSE event (for tool status — non-standard).
	sendEvent := func(event, data string) {
		fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, data)
		flush()
	}

	var fullContent string

	for iteration := 0; iteration < maxToolIterations; iteration++ {
		ch, err := provider.StreamCompletion(c.Request.Context(), cfg, req)
		if err != nil {
			sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(err.Error())))
			return
		}

		var iterContent string
		var toolCalls []providers.ToolCall

		for event := range ch {
			if event.Error != nil {
				sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(event.Error.Error())))
				return
			}

			// Stream text deltas to client
			if event.Delta != "" {
				iterContent += event.Delta
				sendSSE(fmt.Sprintf(
					`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q}`,
					event.Delta, model))
			}

			if event.Done {
				if event.FinishReason == "tool_calls" && len(event.ToolCalls) > 0 {
					toolCalls = event.ToolCalls
				} else {
					// Normal completion — send finish event
					finishReason := event.FinishReason
					if finishReason == "" {
						finishReason = "stop"
					}
					fullContent += iterContent
					sendSSE(fmt.Sprintf(
						`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q}`,
						finishReason, model))
					sendSSE("[DONE]")

					// Persist assistant response
					if fullContent != "" {
						if _, err := h.persistMessage(channelID, userID, "assistant", fullContent, model, 0, 0, nil); err != nil {
							log.Printf("Failed to persist assistant message: %v", err)
						}
					}
					return
				}
				break
			}
		}

		// ── Tool execution round ────────────────
		if len(toolCalls) == 0 {
			// Stream ended without tool calls or finish — shouldn't happen
			sendSSE("[DONE]")
			if iterContent != "" {
				fullContent += iterContent
				if _, err := h.persistMessage(channelID, userID, "assistant", fullContent, model, 0, 0, nil); err != nil {
					log.Printf("Failed to persist assistant message: %v", err)
				}
			}
			return
		}

		// Notify client about tool calls
		toolCallsJSON, _ := json.Marshal(toolCalls)
		sendEvent("tool_use", string(toolCallsJSON))

		// Append assistant message (with tool_calls, possibly with text) to conversation
		assistantMsg := providers.Message{
			Role:    "assistant",
			Content: iterContent,
		}
		for _, tc := range toolCalls {
			assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
		}
		req.Messages = append(req.Messages, assistantMsg)

		// Execute all tool calls
		execCtx := tools.ExecutionContext{
			UserID:    userID,
			ChannelID: channelID,
		}
		for _, tc := range toolCalls {
			call := tools.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			}

			log.Printf("🔧 Executing tool: %s (call %s)", call.Name, call.ID)
			result := tools.ExecuteCall(c.Request.Context(), execCtx, call)

			// Notify client about tool result
			resultJSON, _ := json.Marshal(map[string]interface{}{
				"tool_call_id": result.ToolCallID,
				"name":         result.Name,
				"content":      result.Content,
				"is_error":     result.IsError,
			})
			sendEvent("tool_result", string(resultJSON))

			// Append tool result to conversation for next iteration
			req.Messages = append(req.Messages, providers.Message{
				Role:       "tool",
				Content:    result.Content,
				ToolCallID: result.ToolCallID,
				Name:       result.Name,
			})
		}

		fullContent += iterContent
		// Loop: send updated messages back to provider
	}

	// If we hit max iterations, send what we have
	log.Printf("⚠️ Tool loop hit max iterations (%d) for channel %s", maxToolIterations, channelID)
	sendSSE(fmt.Sprintf(
		`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q}`,
		escapeJSON("[Tool execution limit reached]"), model))
	sendSSE("[DONE]")

	if fullContent != "" {
		if _, err := h.persistMessage(channelID, userID, "assistant", fullContent, model, 0, 0, nil); err != nil {
			log.Printf("Failed to persist assistant message: %v", err)
		}
	}
}

// ── Non-Streaming Completion with Tool Loop ──

func (h *CompletionHandler) syncCompletion(
	c *gin.Context,
	provider providers.Provider,
	cfg providers.ProviderConfig,
	req providers.CompletionRequest,
	channelID, userID, model string,
) {
	var totalInput, totalOutput int
	var finalContent string

	for iteration := 0; iteration < maxToolIterations; iteration++ {
		resp, err := provider.ChatCompletion(c.Request.Context(), cfg, req)
		if err != nil {
			c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()})
			return
		}

		totalInput += resp.InputTokens
		totalOutput += resp.OutputTokens

		// No tool calls — normal response
		if len(resp.ToolCalls) == 0 || resp.FinishReason != "tool_calls" {
			finalContent = resp.Content

			// Persist assistant response
			if _, err := h.persistMessage(channelID, userID, "assistant", finalContent, model, totalInput, totalOutput, nil); err != nil {
				log.Printf("Failed to persist assistant message: %v", err)
			}

			// Return OpenAI-compatible response
			c.JSON(http.StatusOK, gin.H{
				"model": resp.Model,
				"choices": []gin.H{
					{
						"message": gin.H{
							"role":    "assistant",
							"content": finalContent,
						},
						"finish_reason": resp.FinishReason,
					},
				},
				"usage": gin.H{
					"prompt_tokens":     totalInput,
					"completion_tokens": totalOutput,
					"total_tokens":      totalInput + totalOutput,
				},
			})
			return
		}

		// ── Tool execution round ────────────────
		assistantMsg := providers.Message{
			Role:    "assistant",
			Content: resp.Content,
		}
		for _, tc := range resp.ToolCalls {
			assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
		}
		req.Messages = append(req.Messages, assistantMsg)

		execCtx := tools.ExecutionContext{
			UserID:    userID,
			ChannelID: channelID,
		}
		for _, tc := range resp.ToolCalls {
			call := tools.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			}
			log.Printf("🔧 Executing tool: %s (call %s)", call.Name, call.ID)
			result := tools.ExecuteCall(c.Request.Context(), execCtx, call)

			req.Messages = append(req.Messages, providers.Message{
				Role:       "tool",
				Content:    result.Content,
				ToolCallID: result.ToolCallID,
				Name:       result.Name,
			})
		}
	}

	// Max iterations reached
	c.JSON(http.StatusOK, gin.H{
		"model": model,
		"choices": []gin.H{
			{
				"message": gin.H{
					"role":    "assistant",
					"content": "[Tool execution limit reached]",
				},
				"finish_reason": "stop",
			},
		},
	})
}

// ── Model Capabilities ──────────────────────

// escapeJSON escapes a string for safe embedding in a JSON string value.
func escapeJSON(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	s = strings.ReplaceAll(s, "\n", `\n`)
	s = strings.ReplaceAll(s, "\r", `\r`)
	s = strings.ReplaceAll(s, "\t", `\t`)
	return s
}

// getModelCapabilities looks up capabilities from model_configs DB,
// then overlays with known model defaults and heuristic detection.
func (h *CompletionHandler) getModelCapabilities(model, apiConfigID string) providers.ModelCapabilities {
	// Start with known table or heuristic
	// Start with known model table or heuristics
	caps, found := providers.LookupKnownModel(model)
	if !found {
		caps = providers.InferCapabilities(model)
	}

	// Overlay with DB-stored capabilities (from provider sync or admin edit — authoritative)
	var capsJSON []byte
	err := database.DB.QueryRow(`
		SELECT capabilities FROM model_configs
		WHERE model_id = $1 AND api_config_id = $2
	`, model, apiConfigID).Scan(&capsJSON)

	if err == nil && capsJSON != nil {
		var dbCaps providers.ModelCapabilities
		if err := json.Unmarshal(capsJSON, &dbCaps); err == nil {
			if dbCaps.HasProviderData() {
				// DB has real provider data — use as authoritative, fill gaps
				caps = providers.MergeCapabilities(dbCaps, model)
			}
		}
	}

	return caps
}

// ── Config Resolution ───────────────────────
// Priority: request.api_config_id → chat.api_config_id → user's first active config

func (h *CompletionHandler) resolveConfig(userID string, channelID string, req completionRequest) (providers.ProviderConfig, string, string, string, error) {
	var configID string

	// 1. Explicit config from request
	if req.APIConfigID != "" {
		configID = req.APIConfigID
	}

	// 2. Config from channel
	if configID == "" {
		var channelConfigID *string
		err := database.DB.QueryRow(
			`SELECT api_config_id FROM channels WHERE id = $1`, channelID,
		).Scan(&channelConfigID)
		if err == nil && channelConfigID != nil {
			configID = *channelConfigID
		}
	}

	// 3. User's first active config (personal first, then global)
	if configID == "" {
		err := database.DB.QueryRow(`
			SELECT id FROM api_configs
			WHERE (user_id = $1 OR is_global = true OR user_id IS NULL) AND is_active = true
			ORDER BY user_id NULLS LAST, created_at ASC
			LIMIT 1
		`, userID).Scan(&configID)
		if err != nil {
			return providers.ProviderConfig{}, "", "", "", fmt.Errorf("no API config found — add one at /api-configs")
		}
	}

	// Load the config including custom headers and provider settings
	var providerID, endpoint string
	var apiKey, modelDefault *string
	var customHeadersJSON, providerSettingsJSON []byte
	err := database.DB.QueryRow(`
		SELECT provider, endpoint, api_key_encrypted, model_default, custom_headers, provider_settings
		FROM api_configs
		WHERE id = $1 AND (user_id = $2 OR is_global = true OR user_id IS NULL) AND is_active = true
	`, configID, userID).Scan(&providerID, &endpoint, &apiKey, &modelDefault, &customHeadersJSON, &providerSettingsJSON)

	if err == sql.ErrNoRows {
		return providers.ProviderConfig{}, "", "", "", fmt.Errorf("API config not found or not accessible")
	}
	if err != nil {
		return providers.ProviderConfig{}, "", "", "", fmt.Errorf("failed to load API config: %w", err)
	}

	// Resolve model: request > config default
	model := req.Model
	if model == "" && modelDefault != nil {
		model = *modelDefault
	}
	if model == "" {
		return providers.ProviderConfig{}, "", "", "", fmt.Errorf("no model specified and no default model in config")
	}

	key := ""
	if apiKey != nil {
		key = *apiKey
	}

	// Parse custom headers
	customHeaders := make(map[string]string)
	if customHeadersJSON != nil {
		_ = json.Unmarshal(customHeadersJSON, &customHeaders)
	}

	// Parse provider-specific settings
	providerSettings := make(map[string]interface{})
	if providerSettingsJSON != nil {
		_ = json.Unmarshal(providerSettingsJSON, &providerSettings)
	}

	return providers.ProviderConfig{
		Endpoint:      endpoint,
		APIKey:        key,
		CustomHeaders: customHeaders,
		Settings:      providerSettings,
	}, providerID, model, configID, nil
}

// ── Conversation Loader ─────────────────────

// loadConversation builds the message history for the LLM by walking the
// active path through the message tree. Only the current branch is sent —
// sibling branches are never included in context.
func (h *CompletionHandler) loadConversation(channelID, userID, presetSystemPrompt string) ([]providers.Message, error) {
	// Load system prompt from channel
	var systemPrompt *string
	_ = database.DB.QueryRow(
		`SELECT system_prompt FROM channels WHERE id = $1`, channelID,
	).Scan(&systemPrompt)

	messages := make([]providers.Message, 0)

	// Preset system prompt takes priority; channel system prompt is fallback
	if presetSystemPrompt != "" {
		messages = append(messages, providers.Message{
			Role:    "system",
			Content: presetSystemPrompt,
		})
	} else if systemPrompt != nil && *systemPrompt != "" {
		messages = append(messages, providers.Message{
			Role:    "system",
			Content: *systemPrompt,
		})
	}

	// Walk the active path (root → leaf) instead of loading all messages
	path, err := getActivePath(channelID, userID)
	if err != nil {
		return nil, err
	}

	for _, m := range path {
		if m.Role == "system" {
			continue // system prompts handled above
		}
		messages = append(messages, providers.Message{
			Role:    m.Role,
			Content: m.Content,
		})
	}

	return messages, nil
}

// ── Message Persistence ─────────────────────

// persistMessage inserts a message into the tree, updates the cursor, and
// returns the new message's ID.
//
// Parent resolution:
//   - parentOverride (explicit parent for edit/regen operations)
//   - cursor's active_leaf_id (normal conversation flow)
//   - latest message by created_at (fallback for legacy/missing cursor)
func (h *CompletionHandler) persistMessage(channelID, userID, role, content, model string, inputTokens, outputTokens int, parentOverride *string) (string, error) {
	var tokensUsed *int
	if inputTokens > 0 || outputTokens > 0 {
		total := inputTokens + outputTokens
		tokensUsed = &total
	}

	// Determine parent
	var parentID *string
	if parentOverride != nil {
		parentID = parentOverride
	} else {
		parentID, _ = getActiveLeaf(channelID, userID)
	}

	// Compute sibling_index
	siblingIdx := nextSiblingIndex(channelID, parentID)

	// Determine participant
	participantType := "user"
	participantID := userID
	if role == "assistant" {
		participantType = "model"
		participantID = model
	}

	// Insert with RETURNING id
	var newID string
	err := database.DB.QueryRow(`
		INSERT INTO messages (channel_id, role, content, model, tokens_used,
		                      parent_id, participant_type, participant_id, sibling_index)
		VALUES ($1, $2, $3, NULLIF($4, ''), $5, $6, $7, $8, $9)
		RETURNING id
	`, channelID, role, content, model, tokensUsed,
		parentID, participantType, participantID, siblingIdx,
	).Scan(&newID)
	if err != nil {
		return "", err
	}

	// Update cursor to point at new message
	if userID != "" {
		_ = updateCursor(channelID, userID, newID)
	}

	// Touch channel updated_at
	_, _ = database.DB.Exec(`UPDATE channels SET updated_at = NOW() WHERE id = $1`, channelID)
	return newID, nil
}