rony-chat-bot/internal/server/handlers.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/VictorVargas/rony-chat-bot/internal/agent"
	"github.com/VictorVargas/rony-chat-bot/internal/config"
	"github.com/VictorVargas/rony-chat-bot/internal/portfolio"
	"github.com/VictorVargas/rony-chat-bot/internal/streaming"
)

type Handlers struct {
	cfg     *config.Config
	runner  *agent.Runner
	store   *portfolio.Store
	version string
}

func NewHandlers(cfg *config.Config, runner *agent.Runner, store *portfolio.Store, version string) *Handlers {
	return &Handlers{cfg: cfg, runner: runner, store: store, version: version}
}

type ChatRequest struct {
	Messages       []ChatMessage `json:"messages"`
	Stream         *bool         `json:"stream,omitempty"`
	// ConversationID is optional. If empty, the server creates a new
	// conversation and returns its ID in the response (or in the SSE
	// `start` event). Pass an existing ID to continue a previous thread.
	ConversationID string `json:"conversation_id,omitempty"`
}

type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ChatResponse struct {
	ConversationID string         `json:"conversation_id"`
	Content        string         `json:"content"`
	Sources        []string       `json:"sources,omitempty"`
	Usage          streaming.Usage `json:"usage"`
}

func (h *Handlers) Chat(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var req ChatRequest
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&req); err != nil {
		http.Error(w, "invalid JSON body: "+err.Error(), http.StatusBadRequest)
		return
	}
	if err := req.validate(); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	stream := true
	if req.Stream != nil {
		stream = *req.Stream
	}

	// Resolve or create the conversation. If the client passed a non-empty
	// ID we use it as-is; otherwise we mint a new one.
	ctx := r.Context()
	convID := req.ConversationID
	if convID == "" {
		var err error
		convID, err = h.store.CreateConversation(ctx)
		if err != nil {
			http.Error(w, "create conversation: "+err.Error(), http.StatusInternalServerError)
			return
		}
	}

	// Persist the incoming user message. Even if the LLM fails afterward
	// the user sees their question in the conversation history.
	if userMsg := lastUserMessage(req.Messages); userMsg != "" {
		if err := h.store.SaveMessage(ctx, convID, "user", userMsg, nil); err != nil {
			slog.Error("save user message", "err", err)
		}
	}

	history := toAgentMessages(req.Messages)
	if stream {
		h.streamChat(w, r, history, convID)
		return
	}
	h.completeChat(w, r, history, convID)
}

func lastUserMessage(msgs []ChatMessage) string {
	for i := len(msgs) - 1; i >= 0; i-- {
		if msgs[i].Role == "user" {
			return msgs[i].Content
		}
	}
	return ""
}

func toAgentMessages(in []ChatMessage) []agent.Message {
	out := make([]agent.Message, len(in))
	for i, m := range in {
		out[i] = agent.Message{
			Role:    agent.Role(m.Role),
			Content: m.Content,
		}
	}
	return out
}

func (req *ChatRequest) validate() error {
	if len(req.Messages) == 0 {
		return errors.New("messages must not be empty")
	}
	for i, m := range req.Messages {
		switch m.Role {
		case "user", "assistant", "system":
		default:
			return fmt.Errorf("messages[%d].role %q is invalid", i, m.Role)
		}
		if strings.TrimSpace(m.Content) == "" {
			return fmt.Errorf("messages[%d].content is empty", i)
		}
	}
	return nil
}

// streamChat drives the LLM agent and forwards each chunk to the SSE
// stream using the protocol described in docs/architecture.md §3.2.
// convID is the persistent conversation ID — passed in from Chat() after
// resolve-or-create. The assistant reply is saved on success.
func (h *Handlers) streamChat(w http.ResponseWriter, r *http.Request, history []agent.Message, convID string) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}

	ctx := r.Context()
	if err := streaming.WriteStart(w, convID); err != nil {
		slog.Error("sse start failed", "err", err)
		return
	}
	flusher.Flush()

	// Report which RAG sources were used (search happens in BuildMessages).
	_, ragContext, err := h.runner.BuildMessages(ctx, history)
	if err != nil {
		_ = streaming.WriteError(w, "rag: "+err.Error())
		return
	}
	var sources []string
	if ragContext != "" {
		sources = extractSources(ragContext)
		_ = streaming.WriteSources(w, sources)
	}

	// Stream from the model, collecting the full reply as we go.
	var full strings.Builder
	for chunk, err := range h.runner.Stream(ctx, history) {
		if err != nil {
			slog.Error("llm stream", "err", err)
			_ = streaming.WriteError(w, "llm: "+err.Error())
			return
		}
		if chunk.Delta != "" {
			full.WriteString(chunk.Delta)
			if err := streaming.WriteChunk(w, chunk.Delta); err != nil {
				return
			}
		}
		if chunk.FinishReason != "" && chunk.Usage.TotalTokens > 0 {
			h.persistAssistant(ctx, convID, full.String(), sources)
			_ = streaming.WriteDone(w, streaming.Usage{
				InputTokens:  chunk.Usage.InputTokens,
				OutputTokens: chunk.Usage.OutputTokens,
			})
			return
		}
	}
	// Final usage chunk may come on the last iteration; if we never saw a
	// finish_reason + usage in-band, surface what we recorded.
	usage := h.runner.LastUsage()
	h.persistAssistant(ctx, convID, full.String(), sources)
	_ = streaming.WriteDone(w, streaming.Usage{
		InputTokens:  usage.InputTokens,
		OutputTokens: usage.OutputTokens,
	})
}

// persistAssistant saves the full assistant reply. Best-effort: a failure
// here doesn't fail the user's request (the response is already streamed).
func (h *Handlers) persistAssistant(ctx context.Context, convID, content string, sources []string) {
	if strings.TrimSpace(content) == "" {
		return
	}
	if err := h.store.SaveMessage(ctx, convID, "assistant", content, sources); err != nil {
		slog.Error("save assistant message", "err", err, "conv", convID)
	}
}

func (h *Handlers) completeChat(w http.ResponseWriter, r *http.Request, history []agent.Message, convID string) {
	w.Header().Set("Content-Type", "application/json")
	ctx := r.Context()
	var full strings.Builder
	var sources []string
	for chunk, err := range h.runner.Stream(ctx, history) {
		if err != nil {
			http.Error(w, "llm: "+err.Error(), http.StatusBadGateway)
			return
		}
		full.WriteString(chunk.Delta)
	}
	usage := h.runner.LastUsage()
	_, ragContext, _ := h.runner.BuildMessages(ctx, history)
	if ragContext != "" {
		sources = extractSources(ragContext)
	}
	h.persistAssistant(ctx, convID, full.String(), sources)
	resp := ChatResponse{
		ConversationID: convID,
		Content:        full.String(),
		Usage: streaming.Usage{
			InputTokens:  usage.InputTokens,
			OutputTokens: usage.OutputTokens,
		},
		Sources: sources,
	}
	_ = json.NewEncoder(w).Encode(resp)
}

func extractSources(ragContext string) []string {
	// The formatted hits look like "### [N] projectID — section\n..."
	// so we scan the lines for that prefix and dedupe by projectID.
	seen := map[string]bool{}
	var out []string
	for _, line := range strings.Split(ragContext, "\n") {
		if !strings.HasPrefix(line, "### [") {
			continue
		}
		// between "### [N] " and " — "
		rest := strings.TrimPrefix(line, "### [")
		rest = rest[strings.Index(rest, " ")+1:]
		project := rest
		if i := strings.Index(rest, " — "); i > 0 {
			project = rest[:i]
		}
		if !seen[project] {
			seen[project] = true
			out = append(out, project)
		}
	}
	return out
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}

func (h *Handlers) Info(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	p := h.cfg.DefaultProvider()
	_ = json.NewEncoder(w).Encode(map[string]any{
		"name":     h.cfg.Persona.Name,
		"version":  h.version,
		"provider": p.Type,
		"model":    firstNonEmpty(p.Model, p.ModelPath),
		"rag":      h.cfg.RAG.Enabled,
		"top_k":    h.cfg.RAG.TopK,
		"started":  time.Now().UTC().Format(time.RFC3339),
	})
}

func (h *Handlers) Reindex(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	start := time.Now()
	dbPath := h.cfg.RAG.DBPath
	if dbPath == "" {
		http.Error(w, "rag.db_path not configured", http.StatusBadRequest)
		return
	}
	dur, files, chunks, err := portfolio.ReindexOnDisk(dbPath, h.cfg.RAG.DataPath, portfolio.DefaultChunkerConfig())
	if err != nil {
		http.Error(w, "reindex failed: "+err.Error(), http.StatusInternalServerError)
		return
	}
	_ = dur
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"indexed_files": files,
		"total_chunks":  chunks,
		"duration_ms":   time.Since(start).Milliseconds(),
		"db_path":       dbPath,
	})
}

// ---- Conversation REST endpoints -------------------------------------------

// GetConversation returns the full history of one conversation. The
// conversation ID is treated as a bearer token: anyone who knows it can
// read the history. For a public bot this is fine; for private contexts
// add auth at the proxy layer.
func (h *Handlers) GetConversation(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	id := conversationIDFromPath(r.URL.Path)
	if id == "" {
		http.Error(w, "missing conversation id", http.StatusBadRequest)
		return
	}
	conv, err := h.store.GetConversation(r.Context(), id)
	if errors.Is(err, portfolio.ErrConversationNotFound) {
		http.Error(w, "conversation not found", http.StatusNotFound)
		return
	}
	if err != nil {
		http.Error(w, "get conversation: "+err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(conv)
}

// ListConversations returns the most recent N conversation summaries.
// Useful for a "show my chats" sidebar in the widget or a custom UI.
func (h *Handlers) ListConversations(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	limit := 50
	if s := r.URL.Query().Get("limit"); s != "" {
		if n, err := strconv.Atoi(s); err == nil && n > 0 && n <= 200 {
			limit = n
		}
	}
	items, err := h.store.ListConversations(r.Context(), limit)
	if err != nil {
		http.Error(w, "list conversations: "+err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	if items == nil {
		items = []portfolio.ConversationSummary{}
	}
	_ = json.NewEncoder(w).Encode(map[string]any{
		"conversations": items,
		"count":         len(items),
	})
}

// DeleteConversation removes a conversation and all its messages.
func (h *Handlers) DeleteConversation(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodDelete {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	id := conversationIDFromPath(r.URL.Path)
	if id == "" {
		http.Error(w, "missing conversation id", http.StatusBadRequest)
		return
	}
	err := h.store.DeleteConversation(r.Context(), id)
	if errors.Is(err, portfolio.ErrConversationNotFound) {
		http.Error(w, "conversation not found", http.StatusNotFound)
		return
	}
	if err != nil {
		http.Error(w, "delete conversation: "+err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// conversationIDFromPath pulls the ID out of "/api/conversations/{id}".
// chi does this with chi.URLParam(r, "id"); for clarity (and so this
// handler works even without chi) we do it by hand.
func conversationIDFromPath(path string) string {
	const prefix = "/api/conversations/"
	if !strings.HasPrefix(path, prefix) {
		return ""
	}
	id := strings.TrimPrefix(path, prefix)
	// strip trailing slash and any further segments
	if i := strings.IndexByte(id, '/'); i >= 0 {
		id = id[:i]
	}
	return id
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if v != "" {
			return v
		}
	}
	return ""
}