MuyueWorkspace/internal/api/handlers_chat.go

package api

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"regexp"
	"strings"

	"github.com/muyue/muyue/internal/orchestrator"
)

var toolCallRegex = regexp.MustCompile(`\[TOOL_CALL:\{[^\}]+\}\]`)

func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}
	var body struct {
		Message string `json:"message"`
		Stream  bool   `json:"stream"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}
	if body.Message == "" {
		writeError(w, "no message", http.StatusBadRequest)
		return
	}

	s.convStore.Add("user", body.Message)

	if s.convStore.NeedsSummarization() {
		s.autoSummarize()
	}

	orb, err := orchestrator.New(s.config)
	if err != nil {
		writeError(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	orb.SetSystemPrompt(`Tu es l'assistant IA de Muyue Studio. Tu as accès à un outil "crush" pour exécuter des tâches complexes sur l'ordinateur de l'utilisateur.

RÈGLES ABSOLUES:
1. Tu as DEUX possibilités ONLY:
   - Répondre directement à l'utilisateur avec tes connaissances
   - Demander l'exécution d'une tâche via crush en utilisant ce format EXACT:
     [TOOL_CALL:{"tool":"crush","task":"description de la tâche"}]

2. Quand tu utilises [TOOL_CALL:...], le système exécutera la tâche et te donnera le résultat.
   Tu peux ensuite répondre à l'utilisateur avec ce résultat.

3. SOIS CONCIS - pas de blabla, vais droit au but.

4. L'utilisateur ne voit PAS tes pensées entre <think> tags.

5. EXEMPLES d'utilisation de tool:
   - "cherche tous les fichiers .md dans le projet" → [TOOL_CALL:{"tool":"crush","task":"Recherche les fichiers .md dans le projet courant"}]
   - "aide-moi à déboguer cette erreur" → tu peux répondre directement si tu as assez d'info, sinon utiliser tool
   - "quelle est la météo?" → [TOOL_CALL:{"tool":"crush","task":"Cherche la météo actuelle"}]

6. Ne fais PAS de multi-step tool calls dans une seule réponse. Attends le résultat avant de continuer.`)

	if body.Stream {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.WriteHeader(http.StatusOK)
		flusher, canFlush := w.(http.Flusher)

		result, err := orb.SendStream(body.Message, func(chunk string) {
			if strings.HasPrefix(chunk, "<think") {
				data, _ := json.Marshal(map[string]string{"thinking": strings.TrimPrefix(chunk, "<think")})
				w.Write([]byte("data: " + string(data) + "\n\n"))
				if canFlush {
					flusher.Flush()
				}
				return
			}
			if chunk == "</think>" {
				data, _ := json.Marshal(map[string]string{"thinking_end": "true"})
				w.Write([]byte("data: " + string(data) + "\n\n"))
				if canFlush {
					flusher.Flush()
				}
				return
			}
			data, _ := json.Marshal(map[string]string{"content": chunk})
			w.Write([]byte("data: " + string(data) + "\n\n"))
			if canFlush {
				flusher.Flush()
			}
		})
		if err != nil {
			data, _ := json.Marshal(map[string]string{"error": err.Error()})
			w.Write([]byte("data: " + string(data) + "\n\n"))
			if canFlush {
				flusher.Flush()
			}
			return
		}

		// Process tool calls if any
		cleanResult := processToolCalls(result)
		s.convStore.Add("assistant", cleanResult)

		data, _ := json.Marshal(map[string]string{"done": "true"})
		w.Write([]byte("data: " + string(data) + "\n\n"))
		if canFlush {
			flusher.Flush()
		}
		return
	}

	result, err := orb.Send(body.Message)
	if err != nil {
		writeError(w, err.Error(), http.StatusInternalServerError)
		return
	}
	cleanResult := processToolCalls(result)
	s.convStore.Add("assistant", cleanResult)
	writeJSON(w, map[string]string{"content": cleanResult})
}

func processToolCalls(content string) string {
	matches := toolCallRegex.FindAllString(content, -1)
	if len(matches) == 0 {
		return cleanThinkingTags(content)
	}

	var result strings.Builder
	clean := content

	for _, match := range matches {
		// Extract tool and task from [TOOL_CALL:{...}]
		inner := strings.TrimPrefix(match, "[TOOL_CALL:")
		inner = strings.TrimSuffix(inner, "]}") + "}"

		var call struct {
			Tool string `json:"tool"`
			Task string `json:"task"`
		}
		if err := json.Unmarshal([]byte(inner), &call); err != nil {
			continue
		}

		if call.Tool == "crush" && call.Task != "" {
			result.WriteString(fmt.Sprintf("> %s\n\n", call.Task))
			output := executeCrush(call.Task)
			result.WriteString(output)
			result.WriteString("\n\n---\n\n")
		}

		clean = strings.Replace(clean, match, "", 1)
	}

	clean = cleanThinkingTags(clean)

	if result.Len() > 0 {
		clean = strings.TrimSpace(clean) + "\n\n" + strings.TrimSpace(result.String())
	}

	return clean
}

func cleanThinkingTags(content string) string {
	re := regexp.MustCompile(`(?s)<think[^>]*>.*?</think>`)
	return re.ReplaceAllString(content, "")
}

func executeCrush(task string) string {
	cmd := exec.Command("crush", "run", task)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Sprintf("Erreur: %v\n%s", err, string(output))
	}
	return string(output)
}

func (s *Server) autoSummarize() {
	messages := s.convStore.Get()
	if len(messages) < 10 {
		return
	}

	half := len(messages) / 2
	var oldText string
	for _, m := range messages[:half] {
		oldText += m.Role + ": " + m.Content + "\n\n"
	}

	summary := s.convStore.GetSummary()
	if summary != "" {
		oldText = "Résumé précédent:\n" + summary + "\n\nNouveaux échanges:\n" + oldText
	}

	orb, err := orchestrator.New(s.config)
	if err != nil {
		return
	}
	orb.SetSystemPrompt(summarizePrompt)

	result, err := orb.Send(oldText)
	if err != nil {
		return
	}

	s.convStore.SetSummary(result)
	s.convStore.TrimOld(len(messages) - half)
	s.convStore.Add("system", "[Conversation résumée automatiquement]")
}

func (s *Server) handleChatHistory(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeError(w, "GET only", http.StatusMethodNotAllowed)
		return
	}
	messages := s.convStore.Get()
	writeJSON(w, map[string]interface{}{
		"messages": messages,
		"tokens":   s.convStore.ApproxTokenCount(),
	})
}

func (s *Server) handleChatClear(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}
	s.convStore.Clear()
	writeJSON(w, map[string]string{"status": "ok"})
}