MuyueWorkspace/internal/api/handlers_shell_chat.go

package api

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
	"unicode/utf8"

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

type ShellChatRequest struct {
	Message  string `json:"message"`
	Context  string `json:"context,omitempty"`
	Cwd      string `json:"cwd,omitempty"`
	Platform string `json:"platform,omitempty"`
	Stream   bool   `json:"stream"`
}

func (s *Server) handleShellChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}

	if s.shellConvStore.AtLimit() {
		writeError(w, "context limit reached, use /clear", http.StatusBadRequest)
		return
	}

	var req ShellChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}

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

	s.shellConvStore.Add("user", req.Message)

	orb, err := orchestrator.New(s.config)
	if err != nil {
		writeError(w, err.Error(), http.StatusServiceUnavailable)
		return
	}

	orb.SetSystemPrompt(s.buildShellSystemPrompt(req))
	orb.SetTools(s.shellAgentToolsJSON)

	if req.Stream {
		s.handleShellChatStream(w, orb)
	} else {
		s.handleShellChatNonStream(w, orb)
	}
}

func (s *Server) buildShellSystemPrompt(req ShellChatRequest) string {
	var sb strings.Builder

	sb.WriteString(shellSystemPromptBase)

	analysis := LoadSystemAnalysis()
	if analysis != "" {
		sb.WriteString("<system_context>\n")
		sb.WriteString(analysis)
		sb.WriteString("\n</system_context>\n\n")
	}

	sb.WriteString(fmt.Sprintf("OS: %s/%s\n", runtime.GOOS, runtime.GOARCH))
	if hostname, err := os.Hostname(); err == nil {
		sb.WriteString("Hostname: " + hostname + "\n")
	}
	if user := os.Getenv("USER"); user != "" {
		sb.WriteString("User: " + user + "\n")
	}

	canSudo := !agent.NeedsSudoPassword()
	sb.WriteString(fmt.Sprintf("Root: %t\n", !canSudo))
	if canSudo {
		sb.WriteString("⚠️ Session avec privilèges sudo sans mot de passe — les commandes sudo s'exécuteront directement.\n")
	} else {
		sb.WriteString("⚠️ Session sans sudo sans mot de passe — les commandes sudo/doas nécessitent une autorisation. N'utilise PAS sudo ou doas sans demander.\n")
	}

	now := time.Now()
	sb.WriteString(fmt.Sprintf("Date: %s\nHeure: %s\n", now.Format("02/01/2006"), now.Format("15:04:05")))

	return sb.String()
}

func (s *Server) handleShellChatStream(w http.ResponseWriter, orb *orchestrator.Orchestrator) {
	SetupSSEHeaders(w)
	flusher, canFlush := w.(http.Flusher)
	sseWriter := NewSSEWriter(w)

	ctx := context.Background()
	messages := s.buildShellContextMessages()

	engine := NewChatEngine(orb, s.shellAgentRegistry, s.shellAgentToolsJSON)
	engine.SetLimiter(s.AcquireAgentSlot)
	engine.OnChunk(func(data map[string]interface{}) {
		if data == nil {
			return
		}
		sseWriter.Write(data)
		if canFlush {
			flusher.Flush()
		}
	})

	finalContent, allToolCalls, allToolResults, err := engine.RunWithTools(ctx, messages)
	if err != nil {
		sseWriter.Write(map[string]interface{}{"error": err.Error()})
		return
	}

	storeContent := finalContent
	if len(allToolCalls) > 0 {
		storeObj := map[string]interface{}{
			"content":      storeContent,
			"tool_calls":   allToolCalls,
			"tool_results": allToolResults,
		}
		storeJSON, _ := json.Marshal(storeObj)
		storeContent = string(storeJSON)
	}
	s.shellConvStore.Add("assistant", storeContent)

	s.consumption.Record(engine.ProviderName(), engine.TotalTokens)

	sseWriter.Write(map[string]interface{}{
		"done":   "true",
		"tokens": s.shellConvStore.ApproxTokens(),
	})
}

func (s *Server) handleShellChatNonStream(w http.ResponseWriter, orb *orchestrator.Orchestrator) {
	ctx := context.Background()
	messages := s.buildShellContextMessages()

	engine := NewChatEngine(orb, s.shellAgentRegistry, s.shellAgentToolsJSON)
	engine.SetLimiter(s.AcquireAgentSlot)
	finalContent, err := engine.RunNonStream(ctx, messages)
	if err != nil {
		writeError(w, err.Error(), http.StatusInternalServerError)
		return
	}

	s.shellConvStore.Add("assistant", finalContent)

	s.consumption.Record(engine.ProviderName(), engine.TotalTokens)

	writeJSON(w, map[string]interface{}{
		"content": finalContent,
		"tokens":  s.shellConvStore.ApproxTokens(),
	})
}

func (s *Server) buildShellContextMessages() []orchestrator.Message {
	history := s.shellConvStore.Get()

	sysTokens := utf8.RuneCountInString(shellSystemPromptBase) / charsPerToken
	if analysis := LoadSystemAnalysis(); analysis != "" {
		sysTokens += utf8.RuneCountInString(analysis) / charsPerToken
	}
	sysTokens += 100
	toolsTokens := utf8.RuneCountInString(string(s.shellAgentToolsJSON)) / charsPerToken
	responseMargin := 4000

	overhead := sysTokens + toolsTokens + responseMargin
	available := shellMaxTokens - overhead
	if available < 1000 {
		available = 1000
	}

	included := 0
	tokensUsed := 0
	for i := len(history) - 1; i >= 0; i-- {
		displayContent := extractDisplayContent(history[i].Role, history[i].Content)
		msgTokens := utf8.RuneCountInString(displayContent) / charsPerToken
		if msgTokens == 0 {
			msgTokens = 1
		}
		if tokensUsed+msgTokens > available {
			break
		}
		tokensUsed += msgTokens
		included++
	}

	start := len(history) - included
	if start < 0 {
		start = 0
	}

	if start > 0 {
		_ = start
	}

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

	for _, m := range history[start:] {
		if m.Role == "system" {
			continue
		}
		displayContent := extractDisplayContent(m.Role, m.Content)
		messages = append(messages, orchestrator.Message{
			Role:    m.Role,
			Content: orchestrator.TextContent(displayContent),
		})
	}

	return messages
}

func (s *Server) handleShellChatHistory(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeError(w, "GET only", http.StatusMethodNotAllowed)
		return
	}
	messages := s.shellConvStore.Get()
	writeJSON(w, map[string]interface{}{
		"messages":   messages,
		"tokens":     s.shellConvStore.ApproxTokens(),
		"max_tokens": shellMaxTokens,
		"at_limit":   s.shellConvStore.AtLimit(),
	})
}

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

func (s *Server) handleShellAnalyze(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}

	var sysInfo strings.Builder
	sysInfo.WriteString("=== INFORMATIONS SYSTÈME ===\n")
	sysInfo.WriteString(fmt.Sprintf("OS: %s/%s\n", runtime.GOOS, runtime.GOARCH))
	if hostname, err := os.Hostname(); err == nil {
		sysInfo.WriteString("Hostname: " + hostname + "\n")
	}
	if user := os.Getenv("USER"); user != "" {
		sysInfo.WriteString("User: " + user + "\n")
	}

	if data, err := os.ReadFile("/proc/cpuinfo"); err == nil {
		for _, line := range strings.Split(string(data), "\n") {
			if strings.HasPrefix(line, "model name") {
				sysInfo.WriteString("CPU: " + strings.SplitN(line, ":", 2)[1] + "\n")
				break
			}
		}
	}

	if data, err := os.ReadFile("/proc/meminfo"); err == nil {
		for _, line := range strings.Split(string(data), "\n") {
			if strings.HasPrefix(line, "MemTotal:") || strings.HasPrefix(line, "MemAvailable:") {
				sysInfo.WriteString(strings.TrimSpace(line) + "\n")
			}
		}
	}

	if out, err := exec.Command("df", "-h", "/").Output(); err == nil {
		lines := strings.Split(string(out), "\n")
		if len(lines) >= 2 {
			sysInfo.WriteString("Disk: " + strings.TrimSpace(lines[1]) + "\n")
		}
	}

	if out, err := exec.Command("ps", "aux", "--sort=-pcpu").Output(); err == nil {
		lines := strings.Split(string(out), "\n")
		sysInfo.WriteString(fmt.Sprintf("\nProcessus actifs (%d total):\n", len(lines)-1))
		for i := 1; i < len(lines) && i <= 10; i++ {
			fields := strings.Fields(lines[i])
			if len(fields) >= 11 {
				sysInfo.WriteString(fmt.Sprintf("  %-20s CPU:%-6s MEM:%-6s %s\n", fields[10], fields[2]+"%", fields[3]+"%", fields[0]))
			}
		}
	}

	if s.scanResult != nil {
		sysInfo.WriteString("\nOutils installés:\n")
		for _, t := range s.scanResult.Tools {
			status := "✗"
			if t.Installed {
				status = "✓"
			}
			sysInfo.WriteString(fmt.Sprintf("  %s %s %s\n", status, t.Name, t.Version))
		}
	}

	orb, err := orchestrator.New(s.config)
	if err != nil {
		writeError(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	orb.SetSystemPrompt(agent.StudioSystemPrompt())

	analysisPrompt := `Tu es un expert en administration système. Analyse les informations suivantes et génère un rapport structuré en markdown.

STRUCTURE REQUISE :

## État du système
- Résumé en 2-3 phrases de l'état général (OK/Attention/Critique)

## Points d'attention
Liste les problèmes détectés par priorité :
- **CRITIQUE** : problèmes de sécurité, espace disque < 10%, mémoire < 10%
- **ATTENTION** : CPU élevé, services en échec, config non-optimale
- **INFO** : améliorations possibles, mises à jour disponibles

## Recommandations
Pour chaque point d'attention, donne UNE commande ou action corrective concrète.

## Outils manquants
Liste les outils utiles non installés avec la commande d'installation.

## Réseau
- Interfaces actives, ports en écoute, connectivité

RÈGLES :
- Pas de blabla générique — sois spécifique à CE système
- Inclus les valeurs numériques réelles (%, Go, MHz)
- Max 1500 mots
- Le rapport sert de contexte persistant pour un assistant terminal

` + sysInfo.String()

	result, err := orb.Send(analysisPrompt)
	if err != nil {
		writeError(w, "analysis failed: "+err.Error(), http.StatusInternalServerError)
		return
	}

	SaveSystemAnalysis(result)

	writeJSON(w, map[string]interface{}{
		"status":   "ok",
		"analysis": result,
	})
}

func (s *Server) handleShellAnalysisGet(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeError(w, "GET only", http.StatusMethodNotAllowed)
		return
	}
	analysis := LoadSystemAnalysis()
	if analysis == "" {
		writeJSON(w, map[string]interface{}{"analysis": nil})
		return
	}
	writeJSON(w, map[string]interface{}{"analysis": analysis})
}