core/server/handlers/messages.go

package handlers

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

	"github.com/gin-gonic/gin"

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

// ── Request / Response types ────────────────

type createMessageRequest struct {
	Role    string `json:"role" binding:"required,oneof=user assistant system"`
	Content string `json:"content" binding:"required"`
	Model   string `json:"model,omitempty"`
}

type messageResponse struct {
	ID           string  `json:"id"`
	ChannelID    string  `json:"channel_id"`
	Role         string  `json:"role"`
	Content      string  `json:"content"`
	Model        *string `json:"model"`
	TokensUsed   *int    `json:"tokens_used"`
	ParentID     *string `json:"parent_id,omitempty"`
	SiblingCount int     `json:"sibling_count"`
	SiblingIndex int     `json:"sibling_index"`
	CreatedAt    string  `json:"created_at"`
}

type editRequest struct {
	Content string `json:"content" binding:"required"`
}

type regenerateRequest struct {
	Model       string   `json:"model,omitempty"`
	PresetID    string   `json:"preset_id,omitempty"`
	APIConfigID string   `json:"provider_config_id,omitempty"`
	MaxTokens   int      `json:"max_tokens,omitempty"`
	Temperature *float64 `json:"temperature,omitempty"`
}

type cursorRequest struct {
	ActiveLeafID string `json:"active_leaf_id" binding:"required"`
}

// MessageHandler holds dependencies for message endpoints.
type MessageHandler struct{}

// NewMessageHandler creates a new message handler.
func NewMessageHandler() *MessageHandler {
	return &MessageHandler{}
}

// ── List Messages (flat, all branches) ──────
// GET /channels/:id/messages
//
// Returns all live messages across all branches (useful for admin/debug).
// Frontend should prefer /channels/:id/path for rendering.

func (h *MessageHandler) ListMessages(c *gin.Context) {
	page, perPage, offset := parsePagination(c)

	userID := getUserID(c)
	channelID := c.Param("id")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	var total int
	err := database.DB.QueryRow(
		`SELECT COUNT(*) FROM messages WHERE channel_id = $1 AND deleted_at IS NULL`,
		channelID,
	).Scan(&total)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to count messages"})
		return
	}

	rows, err := database.DB.Query(`
		SELECT id, channel_id, role, content, model, tokens_used, parent_id,
		       sibling_index, created_at
		FROM messages
		WHERE channel_id = $1 AND deleted_at IS NULL
		ORDER BY created_at ASC
		LIMIT $2 OFFSET $3
	`, channelID, perPage, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list messages"})
		return
	}
	defer rows.Close()

	messages := make([]messageResponse, 0)
	for rows.Next() {
		var msg messageResponse
		err := rows.Scan(
			&msg.ID, &msg.ChannelID, &msg.Role, &msg.Content,
			&msg.Model, &msg.TokensUsed, &msg.ParentID,
			&msg.SiblingIndex, &msg.CreatedAt,
		)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to scan message"})
			return
		}
		msg.SiblingCount = getSiblingCount(channelID, msg.ParentID)
		messages = append(messages, msg)
	}

	c.JSON(http.StatusOK, paginatedResponse{
		Data:       messages,
		Page:       page,
		PerPage:    perPage,
		Total:      total,
		TotalPages: int(math.Ceil(float64(total) / float64(perPage))),
	})
}

// ── Active Path ─────────────────────────────
// GET /channels/:id/path
//
// Returns the active branch from root → leaf with sibling metadata
// at each node. This is what the frontend renders.

func (h *MessageHandler) GetActivePath(c *gin.Context) {
	userID := getUserID(c)
	channelID := c.Param("id")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	path, err := getActivePath(channelID, userID)
	if err != nil {
		log.Printf("GetActivePath error: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load path"})
		return
	}

	c.JSON(http.StatusOK, gin.H{"path": path})
}

// ── Create Message (manual) ─────────────────
// POST /channels/:id/messages

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

	userID := getUserID(c)
	channelID := c.Param("id")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	// Use cursor for parent, compute sibling_index
	parentID, _ := getActiveLeaf(channelID, userID)
	siblingIdx := nextSiblingIndex(channelID, parentID)

	participantType := "user"
	participantID := userID
	if req.Role == "assistant" {
		participantType = "model"
		if req.Model != "" {
			participantID = req.Model
		}
	}

	var msg messageResponse
	err := database.DB.QueryRow(`
		INSERT INTO messages (channel_id, role, content, model, parent_id,
		                      participant_type, participant_id, sibling_index)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
		RETURNING id, channel_id, role, content, model, tokens_used, parent_id,
		          sibling_index, created_at
	`, channelID, req.Role, req.Content, req.Model, parentID,
		participantType, participantID, siblingIdx,
	).Scan(
		&msg.ID, &msg.ChannelID, &msg.Role, &msg.Content,
		&msg.Model, &msg.TokensUsed, &msg.ParentID,
		&msg.SiblingIndex, &msg.CreatedAt,
	)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create message"})
		return
	}
	msg.SiblingCount = getSiblingCount(channelID, msg.ParentID)

	_ = updateCursor(channelID, userID, msg.ID)
	_, _ = database.DB.Exec(`UPDATE channels SET updated_at = NOW() WHERE id = $1`, channelID)

	c.JSON(http.StatusCreated, msg)
}

// ── Edit Message (create sibling) ───────────
// POST /channels/:id/messages/:msgId/edit
//
// Creates a new user message as a sibling of the target (same parent_id).
// Does NOT auto-trigger completion — the frontend sends a separate request.

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

	userID := getUserID(c)
	channelID := c.Param("id")
	messageID := c.Param("msgId")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	// Load target — must exist, belong to channel, be a user message
	var targetParentID *string
	var targetRole string
	err := database.DB.QueryRow(`
		SELECT parent_id, role FROM messages
		WHERE id = $1 AND channel_id = $2 AND deleted_at IS NULL
	`, messageID, channelID).Scan(&targetParentID, &targetRole)

	if err == sql.ErrNoRows {
		c.JSON(http.StatusNotFound, gin.H{"error": "message not found"})
		return
	}
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load message"})
		return
	}
	if targetRole != "user" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "can only edit user messages"})
		return
	}

	// Create sibling: same parent_id as the target
	siblingIdx := nextSiblingIndex(channelID, targetParentID)

	var msg messageResponse
	err = database.DB.QueryRow(`
		INSERT INTO messages (channel_id, role, content, parent_id,
		                      participant_type, participant_id, sibling_index)
		VALUES ($1, 'user', $2, $3, 'user', $4, $5)
		RETURNING id, channel_id, role, content, model, tokens_used, parent_id,
		          sibling_index, created_at
	`, channelID, req.Content, targetParentID, userID, siblingIdx,
	).Scan(
		&msg.ID, &msg.ChannelID, &msg.Role, &msg.Content,
		&msg.Model, &msg.TokensUsed, &msg.ParentID,
		&msg.SiblingIndex, &msg.CreatedAt,
	)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create edit"})
		return
	}
	msg.SiblingCount = getSiblingCount(channelID, msg.ParentID)

	// Cursor now points to the new sibling (it's a leaf — no children yet)
	_ = updateCursor(channelID, userID, msg.ID)
	_, _ = database.DB.Exec(`UPDATE channels SET updated_at = NOW() WHERE id = $1`, channelID)

	c.JSON(http.StatusCreated, msg)
}

// ── Regenerate / Complete ──────────────────────
// POST /channels/:id/messages/:msgId/regenerate
//
// Two modes depending on target role:
//   - assistant message → creates a NEW assistant sibling (same parent_id).
//     Context is root → target's parent. ("Give me a different response.")
//   - user message → creates a child assistant response.
//     Context is root → target. ("Respond to this message.")
//
// The second mode is used after editing: the edit endpoint creates a
// sibling user message, then the frontend calls regenerate on it to
// get an assistant response without duplicating the user message.

func (h *MessageHandler) Regenerate(c *gin.Context) {
	var req regenerateRequest
	_ = c.ShouldBindJSON(&req) // body is optional

	userID := getUserID(c)
	channelID := c.Param("id")
	messageID := c.Param("msgId")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	// Load target message
	var targetParentID *string
	var targetRole string
	err := database.DB.QueryRow(`
		SELECT parent_id, role FROM messages
		WHERE id = $1 AND channel_id = $2 AND deleted_at IS NULL
	`, messageID, channelID).Scan(&targetParentID, &targetRole)

	if err == sql.ErrNoRows {
		c.JSON(http.StatusNotFound, gin.H{"error": "message not found"})
		return
	}
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load message"})
		return
	}
	if targetRole != "assistant" && targetRole != "user" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "can only regenerate user or assistant messages"})
		return
	}

	// Determine context path and new message's parent based on target role:
	//
	//   assistant target: context = root → target's parent (exclude target)
	//                     new parent = target's parent (sibling of target)
	//
	//   user target:      context = root → target (include target)
	//                     new parent = target itself (child of target)
	var contextPath []PathMessage
	var newParentID *string

	if targetRole == "assistant" {
		if targetParentID == nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "cannot regenerate root message"})
			return
		}
		contextPath, err = getPathToLeaf(channelID, *targetParentID)
		newParentID = targetParentID
	} else {
		// user message — respond to it
		contextPath, err = getPathToLeaf(channelID, messageID)
		newParentID = &messageID
	}

	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to build context"})
		return
	}

	// ── Resolve model + provider ──

	comp := NewCompletionHandler()

	var presetSystemPrompt string
	model := req.Model
	apiConfigID := req.APIConfigID
	temperature := req.Temperature
	maxTokens := req.MaxTokens

	if req.PresetID != "" {
		preset := ResolvePreset(req.PresetID, userID)
		if preset == nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "preset not found"})
			return
		}
		if model == "" {
			model = preset.BaseModelID
		}
		if apiConfigID == "" && preset.ProviderConfigID != nil {
			apiConfigID = *preset.ProviderConfigID
		}
		if temperature == nil && preset.Temperature != nil {
			temperature = preset.Temperature
		}
		if maxTokens == 0 && preset.MaxTokens != nil {
			maxTokens = *preset.MaxTokens
		}
		if preset.SystemPrompt != "" {
			presetSystemPrompt = preset.SystemPrompt
		}
	}

	// Fallback: channel's stored model
	if model == "" {
		var channelModel *string
		_ = database.DB.QueryRow(`SELECT model FROM channels WHERE id = $1`, channelID).Scan(&channelModel)
		if channelModel != nil {
			model = *channelModel
		}
	}

	providerCfg, providerID, model, configID, err := comp.resolveConfig(userID, channelID, completionRequest{
		Model:       model,
		APIConfigID: apiConfigID,
	})
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

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

	// Build LLM message array
	llmMessages := make([]providers.Message, 0, len(contextPath)+1)

	if presetSystemPrompt != "" {
		llmMessages = append(llmMessages, providers.Message{Role: "system", Content: presetSystemPrompt})
	} else {
		var systemPrompt *string
		_ = database.DB.QueryRow(`SELECT system_prompt FROM channels WHERE id = $1`, channelID).Scan(&systemPrompt)
		if systemPrompt != nil && *systemPrompt != "" {
			llmMessages = append(llmMessages, providers.Message{Role: "system", Content: *systemPrompt})
		}
	}

	for _, m := range contextPath {
		if m.Role == "system" {
			continue
		}
		llmMessages = append(llmMessages, providers.Message{Role: m.Role, Content: m.Content})
	}

	provReq := providers.CompletionRequest{
		Model:    model,
		Messages: llmMessages,
	}

	caps := comp.getModelCapabilities(c, model, configID)
	if maxTokens > 0 {
		provReq.MaxTokens = maxTokens
	} else {
		provReq.MaxTokens = capspkg.ResolveMaxOutput(model, caps)
	}
	if temperature != nil {
		provReq.Temperature = temperature
	}

	// ── Stream the response ──

	ch, err := provider.StreamCompletion(c.Request.Context(), providerCfg, provReq)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()})
		return
	}

	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")
	c.Status(http.StatusOK)

	var fullContent string
	flusher, _ := c.Writer.(http.Flusher)

	for event := range ch {
		if event.Error != nil {
			fmt.Fprintf(c.Writer, "data: {\"error\":\"%s\"}\n\n", event.Error.Error())
			if flusher != nil {
				flusher.Flush()
			}
			break
		}

		if event.Delta != "" {
			fullContent += event.Delta
			fmt.Fprintf(c.Writer, "data: {\"choices\":[{\"delta\":{\"content\":%q},\"finish_reason\":null}],\"model\":%q}\n\n",
				event.Delta, model)
			if flusher != nil {
				flusher.Flush()
			}
		}

		if event.Done {
			finishReason := event.FinishReason
			if finishReason == "" {
				finishReason = "stop"
			}
			fmt.Fprintf(c.Writer, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":%q}],\"model\":%q}\n\n",
				finishReason, model)
			io.WriteString(c.Writer, "data: [DONE]\n\n")
			if flusher != nil {
				flusher.Flush()
			}
			break
		}
	}

	// Persist as child or sibling depending on target role
	if fullContent != "" {
		siblingIdx := nextSiblingIndex(channelID, newParentID)

		var newID string
		err := database.DB.QueryRow(`
			INSERT INTO messages (channel_id, role, content, model,
			                      parent_id, participant_type, participant_id, sibling_index)
			VALUES ($1, 'assistant', $2, $3, $4, 'model', $5, $6)
			RETURNING id
		`, channelID, fullContent, model, newParentID, model, siblingIdx).Scan(&newID)

		if err != nil {
			log.Printf("Failed to persist regenerated message: %v", err)
		} else {
			_ = updateCursor(channelID, userID, newID)

			// Send message metadata in a final SSE event so frontend can update
			msgJSON, _ := json.Marshal(gin.H{
				"id":            newID,
				"parent_id":     newParentID,
				"sibling_index": siblingIdx,
				"sibling_count": getSiblingCount(channelID, newParentID),
			})
			fmt.Fprintf(c.Writer, "data: {\"switchboard_meta\":%s}\n\n", msgJSON)
			if flusher != nil {
				flusher.Flush()
			}
		}

		_, _ = database.DB.Exec(`UPDATE channels SET updated_at = NOW() WHERE id = $1`, channelID)
	}
}

// ── Switch Branch (update cursor) ───────────
// PUT /channels/:id/cursor
//
// Navigates to a different branch. If the target message has children,
// follows the first-child chain to find the leaf.

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

	userID := getUserID(c)
	channelID := c.Param("id")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	// Verify the target message belongs to this channel
	var msgChannelID string
	err := database.DB.QueryRow(`
		SELECT channel_id FROM messages
		WHERE id = $1 AND deleted_at IS NULL
	`, req.ActiveLeafID).Scan(&msgChannelID)

	if err == sql.ErrNoRows {
		c.JSON(http.StatusNotFound, gin.H{"error": "message not found"})
		return
	}
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to verify message"})
		return
	}
	if msgChannelID != channelID {
		c.JSON(http.StatusForbidden, gin.H{"error": "message does not belong to this channel"})
		return
	}

	// Walk down to the leaf from the target
	leafID, err := findLeafFromMessage(req.ActiveLeafID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to find leaf"})
		return
	}

	if err := updateCursor(channelID, userID, leafID); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update cursor"})
		return
	}

	// Return the new active path
	path, err := getPathToLeaf(channelID, leafID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load path"})
		return
	}

	c.JSON(http.StatusOK, gin.H{"path": path, "active_leaf_id": leafID})
}

// ── List Siblings ───────────────────────────
// GET /channels/:id/messages/:msgId/siblings
//
// Returns all siblings of a message (messages sharing the same parent_id).
// Used by the branch navigator to populate the arrow navigation.

func (h *MessageHandler) ListSiblings(c *gin.Context) {
	userID := getUserID(c)
	channelID := c.Param("id")
	messageID := c.Param("msgId")

	if !userOwnsChannel(c, channelID, userID) {
		return
	}

	siblings, currentIdx, err := getSiblings(messageID)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "message not found"})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"siblings":      siblings,
		"current_index": currentIdx,
		"total":         len(siblings),
	})
}

// ── Ownership Check ─────────────────────────

func userOwnsChannel(c *gin.Context, channelID, userID string) bool {
	var ownerID string
	err := database.DB.QueryRow(
		`SELECT user_id FROM channels WHERE id = $1`, channelID,
	).Scan(&ownerID)

	if err == sql.ErrNoRows {
		c.JSON(http.StatusNotFound, gin.H{"error": "channel not found"})
		return false
	}
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to verify channel ownership"})
		return false
	}
	if ownerID != userID {
		c.JSON(http.StatusForbidden, gin.H{"error": "not your channel"})
		return false
	}
	return true
}