homelabby/internal/advisor/handler.go

package advisor

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"

	"github.com/go-chi/chi/v5"
	openai "github.com/sashabaranov/go-openai"

	"git.georgsen.dk/hwlab/internal/ai"
	"git.georgsen.dk/hwlab/internal/store"
)

const (
	defaultModel    = "anthropic/claude-opus-4"
	maxMessageBytes = 64 * 1024 // 64 KB body limit (T-06-02-03)
	maxMessageChars = 8000      // truncate message content to 8000 chars (T-06-02-03)
)

// ChatRequest is the JSON body for POST /api/advisor/chat.
type ChatRequest struct {
	ConversationID string `json:"conversation_id"` // empty = new conversation
	Message        string `json:"message"`
	Model          string `json:"model"` // empty = defaultModel
}

// AdvisorHandler implements the three advisor endpoints.
type AdvisorHandler struct {
	store *store.Store
	ctx   *InventoryContextBuilder
	aiCfg ai.AIConfig
}

// NewAdvisorHandler creates an AdvisorHandler.
func NewAdvisorHandler(s *store.Store, ctxBuilder *InventoryContextBuilder, aiCfg ai.AIConfig) *AdvisorHandler {
	return &AdvisorHandler{
		store: s,
		ctx:   ctxBuilder,
		aiCfg: aiCfg,
	}
}

// StreamChat handles POST /api/advisor/chat.
// It streams tokens from the configured AI model back to the client via SSE.
// Each token is sent as: data: {"conversation_id":"...","token":"..."}\n\n
// The final event is: data: [DONE]\n\n
//
// Security:
//   - Body is limited to 64 KB (T-06-02-03)
//   - Message content is truncated to 8000 chars (T-06-02-03)
//   - API key is never written to the SSE stream (T-06-02-02)
func (h *AdvisorHandler) StreamChat(w http.ResponseWriter, r *http.Request) {
	// Enforce body size limit (T-06-02-03).
	r.Body = http.MaxBytesReader(w, r.Body, maxMessageBytes)

	var req ChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request body", http.StatusBadRequest)
		return
	}

	if req.Message == "" {
		http.Error(w, "message is required", http.StatusBadRequest)
		return
	}

	// Truncate message to guard against overly long inputs (T-06-02-03).
	if len(req.Message) > maxMessageChars {
		req.Message = req.Message[:maxMessageChars]
	}

	model := req.Model
	if model == "" {
		model = defaultModel
	}

	ctx := r.Context()

	// Create or reuse conversation.
	convID := req.ConversationID
	if convID == "" {
		id, err := h.store.CreateConversation(ctx, model)
		if err != nil {
			log.Printf("advisor: create conversation: %v", err)
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}
		convID = id
	}

	// Persist user message.
	if _, err := h.store.AddMessage(ctx, convID, "user", req.Message); err != nil {
		log.Printf("advisor: add user message: %v", err)
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}

	// Build inventory context for system prompt.
	invCtx, err := h.ctx.BuildContext(ctx)
	if err != nil {
		log.Printf("advisor: build inventory context: %v (proceeding without inventory)", err)
		invCtx = "Inventory: unavailable\n"
	}

	systemPrompt := "You are a homelab advisor. Here is the current inventory:\n\n" + invCtx

	messages := []openai.ChatCompletionMessage{
		{Role: openai.ChatMessageRoleSystem, Content: systemPrompt},
		{Role: openai.ChatMessageRoleUser, Content: req.Message},
	}

	// Build an OpenAI-compatible client pointed at Tier3 (OpenRouter).
	// API key is read from config only — never from the request (T-06-02-02).
	oCfg := openai.DefaultConfig(h.aiCfg.Tier3.APIKey)
	oCfg.BaseURL = h.aiCfg.Tier3.BaseURL
	client := openai.NewClientWithConfig(oCfg)

	stream, err := client.CreateChatCompletionStream(ctx, openai.ChatCompletionRequest{
		Model:    model,
		Messages: messages,
		Stream:   true,
	})
	if err != nil {
		log.Printf("advisor: create stream: %v", err)
		http.Error(w, "upstream AI error", http.StatusBadGateway)
		return
	}
	defer stream.Close()

	// Set SSE headers.
	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")

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

	type tokenEvent struct {
		ConversationID string `json:"conversation_id"`
		Token          string `json:"token,omitempty"`
	}
	type errorEvent struct {
		Error string `json:"error"`
	}

	var fullContent strings.Builder
	for {
		resp, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			log.Printf("advisor: stream recv: %v", err)
			errPayload, _ := json.Marshal(errorEvent{Error: "stream interrupted"})
			fmt.Fprintf(w, "data: %s\n\n", errPayload)
			flusher.Flush()
			return
		}

		if len(resp.Choices) == 0 {
			continue
		}

		token := resp.Choices[0].Delta.Content
		if token == "" {
			continue
		}
		fullContent.WriteString(token)

		payload, _ := json.Marshal(tokenEvent{ConversationID: convID, Token: token})
		fmt.Fprintf(w, "data: %s\n\n", payload)
		flusher.Flush()
	}

	// Send DONE sentinel.
	fmt.Fprintf(w, "data: [DONE]\n\n")
	flusher.Flush()

	// Persist assistant response.
	if _, err := h.store.AddMessage(ctx, convID, "assistant", fullContent.String()); err != nil {
		log.Printf("advisor: add assistant message: %v", err)
	}
}

// GetConversations handles GET /api/advisor/conversations.
func (h *AdvisorHandler) GetConversations(w http.ResponseWriter, r *http.Request) {
	list, err := h.store.ListConversations(r.Context())
	if err != nil {
		log.Printf("advisor: list conversations: %v", err)
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(list); err != nil {
		log.Printf("advisor: encode conversations: %v", err)
	}
}

// GetConversation handles GET /api/advisor/conversations/{id}.
func (h *AdvisorHandler) GetConversation(w http.ResponseWriter, r *http.Request) {
	id := chi.URLParam(r, "id")
	conv, err := h.store.GetConversation(r.Context(), id)
	if err != nil {
		if errors.Is(err, store.ErrNotFound) {
			http.Error(w, "conversation not found", http.StatusNotFound)
			return
		}
		log.Printf("advisor: get conversation %s: %v", id, err)
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(conv); err != nil {
		log.Printf("advisor: encode conversation: %v", err)
	}
}