MuyueWorkspace/internal/api/conversation.go

package api

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

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

const contextWindowTokens = 150000
const summarizeRatio = 0.80
const charsPerToken = 4

func extractDisplayContent(role, content string) string {
	if role != "assistant" {
		return content
	}
	var parsed struct {
		Content     string `json:"content"`
		ToolCalls   []struct {
			Name string `json:"name"`
			Args string `json:"args"`
		} `json:"tool_calls"`
		ToolResults []struct {
			Name   string `json:"name"`
			Result string `json:"result"`
		} `json:"tool_results"`
	}
	if err := json.Unmarshal([]byte(content), &parsed); err != nil {
		return content
	}
	var sb strings.Builder
	if parsed.Content != "" {
		sb.WriteString(parsed.Content)
	}
	for _, tc := range parsed.ToolCalls {
		sb.WriteString("\n[")
		sb.WriteString(tc.Name)
		sb.WriteString("] ")
		sb.WriteString(tc.Args)
	}
	for _, tr := range parsed.ToolResults {
		sb.WriteString("\n[result")
		if tr.Name != "" {
			sb.WriteString(":")
			sb.WriteString(tr.Name)
		}
		sb.WriteString("] ")
		sb.WriteString(tr.Result)
	}
	return sb.String()
}

type FeedMessage struct {
	ID         string   `json:"id"`
	Role       string   `json:"role"`
	Content    string   `json:"content"`
	Time       string   `json:"time"`
	Images     []string `json:"images,omitempty"`
	Summarized bool     `json:"summarized,omitempty"`
}

type Conversation struct {
	Messages  []FeedMessage `json:"messages"`
	Summary   string        `json:"summary,omitempty"`
	CreatedAt string        `json:"created_at"`
	UpdatedAt string        `json:"updated_at"`
}

type ConversationStore struct {
	mu   sync.RWMutex
	path string
	conv *Conversation
}

type TokenCount struct {
	total     int
	byRole    map[string]int
	byMessage int
}

type SearchResult struct {
	ID      string `json:"id"`
	Role    string `json:"role"`
	Content string `json:"content"`
	Time    string `json:"time"`
}

func NewConversationStore() *ConversationStore {
	dir, err := config.ConfigDir()
	if err != nil {
		dir = "/tmp/muyue"
	}
	path := filepath.Join(dir, "conversation.json")
	cs := &ConversationStore{path: path}
	cs.load()
	return cs
}

func (cs *ConversationStore) load() {
	data, err := os.ReadFile(cs.path)
	if err != nil {
		cs.conv = &Conversation{
			Messages:  []FeedMessage{},
			CreatedAt: time.Now().Format(time.RFC3339),
			UpdatedAt: time.Now().Format(time.RFC3339),
		}
		return
	}
	var conv Conversation
	if err := json.Unmarshal(data, &conv); err != nil {
		cs.conv = &Conversation{
			Messages:  []FeedMessage{},
			CreatedAt: time.Now().Format(time.RFC3339),
			UpdatedAt: time.Now().Format(time.RFC3339),
		}
		return
	}
	if conv.Messages == nil {
		conv.Messages = []FeedMessage{}
	}
	cs.conv = &conv
}

func (cs *ConversationStore) save() error {
	cs.conv.UpdatedAt = time.Now().Format(time.RFC3339)
	data, err := json.MarshalIndent(cs.conv, "", "  ")
	if err != nil {
		return err
	}
	dir := filepath.Dir(cs.path)
	os.MkdirAll(dir, 0755)
	return os.WriteFile(cs.path, data, 0600)
}

func (cs *ConversationStore) Get() []FeedMessage {
	cs.mu.RLock()
	defer cs.mu.RUnlock()
	out := make([]FeedMessage, len(cs.conv.Messages))
	copy(out, cs.conv.Messages)
	return out
}

func (cs *ConversationStore) GetSummary() string {
	cs.mu.RLock()
	defer cs.mu.RUnlock()
	return cs.conv.Summary
}

func (cs *ConversationStore) Add(role, content string) FeedMessage {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	msg := FeedMessage{
		ID:      generateMsgID(),
		Role:    role,
		Content: content,
		Time:    time.Now().Format(time.RFC3339),
	}
	cs.conv.Messages = append(cs.conv.Messages, msg)
	cs.save()
	return msg
}

func (cs *ConversationStore) AddWithImages(role, content string, imageIDs []string) FeedMessage {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	msg := FeedMessage{
		ID:      generateMsgID(),
		Role:    role,
		Content: content,
		Time:    time.Now().Format(time.RFC3339),
		Images:  imageIDs,
	}
	cs.conv.Messages = append(cs.conv.Messages, msg)
	cs.save()
	return msg
}

func (cs *ConversationStore) Clear() {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	var imageIDs []string
	for _, m := range cs.conv.Messages {
		imageIDs = append(imageIDs, m.Images...)
	}

	cs.conv.Messages = []FeedMessage{}
	cs.conv.Summary = ""
	cs.conv.CreatedAt = time.Now().Format(time.RFC3339)
	cs.conv.UpdatedAt = time.Now().Format(time.RFC3339)
	cs.save()

	go cleanupImages(imageIDs)
}

func (cs *ConversationStore) SetSummary(summary string) {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	cs.conv.Summary = summary
	cs.save()
}

func (cs *ConversationStore) MarkSummarized(upToIndex int) {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	if upToIndex <= 0 || upToIndex >= len(cs.conv.Messages) {
		return
	}
	for i := 0; i < upToIndex; i++ {
		cs.conv.Messages[i].Summarized = true
	}
	cs.save()
}

func (cs *ConversationStore) ApproxTokenCount() int {
	return cs.ApproxTokenCountDetailed().total
}

func (cs *ConversationStore) ApproxTokenCountDetailed() TokenCount {
	cs.mu.RLock()
	defer cs.mu.RUnlock()

	result := TokenCount{
		byRole: make(map[string]int),
	}

	for _, m := range cs.conv.Messages {
		if m.Role == "system" || m.Summarized {
			continue
		}
		count := utf8.RuneCountInString(extractDisplayContent(m.Role, m.Content)) / charsPerToken
		result.byMessage += count
		result.byRole[m.Role] += count
	}

	if cs.conv.Summary != "" {
		result.total = result.byMessage + utf8.RuneCountInString(cs.conv.Summary)/charsPerToken
	} else {
		result.total = result.byMessage
	}

	return result
}

func (cs *ConversationStore) NeedsSummarization() bool {
	return cs.ApproxTokenCount() > int(float64(contextWindowTokens)*summarizeRatio)
}

func (cs *ConversationStore) Search(query string) []SearchResult {
	cs.mu.RLock()
	defer cs.mu.RUnlock()

	var results []SearchResult
	queryLower := strings.ToLower(query)

	for _, msg := range cs.conv.Messages {
		if strings.Contains(strings.ToLower(msg.Content), queryLower) {
			results = append(results, SearchResult{
				ID:      msg.ID,
				Role:    msg.Role,
				Content: msg.Content,
				Time:    msg.Time,
			})
		}
	}

	return results
}

func (cs *ConversationStore) ExportMarkdown() string {
	cs.mu.RLock()
	defer cs.mu.RUnlock()

	var sb strings.Builder
	sb.WriteString("# Conversation Export\n\n")
	sb.WriteString(fmt.Sprintf("Exporté le: %s\n\n", time.Now().Format(time.RFC3339)))

	if cs.conv.Summary != "" {
		sb.WriteString("## Résumé\n\n")
		sb.WriteString(cs.conv.Summary)
		sb.WriteString("\n\n---\n\n")
	}

	sb.WriteString("## Messages\n\n")

	for i, msg := range cs.conv.Messages {
		roleLabel := msg.Role
		if roleLabel == "user" {
			roleLabel = "👤 Utilisateur"
		} else if roleLabel == "assistant" {
			roleLabel = "🤖 Assistant"
		} else if roleLabel == "system" {
			roleLabel = "⚙️ Système"
		}

		timestamp := ""
		if msg.Time != "" {
			if t, err := time.Parse(time.RFC3339, msg.Time); err == nil {
				timestamp = t.Format("2006-01-02 15:04")
			}
		}

		sb.WriteString(fmt.Sprintf("### [%d] %s (%s)\n\n", i+1, roleLabel, timestamp))
		sb.WriteString(msg.Content)
		sb.WriteString("\n\n---\n\n")
	}

	return sb.String()
}

func (cs *ConversationStore) ExportJSON() string {
	cs.mu.RLock()
	defer cs.mu.RUnlock()

	data, err := json.MarshalIndent(cs.conv, "", "  ")
	if err != nil {
		return "{}"
	}
	return string(data)
}

func generateMsgID() string {
	return time.Now().Format("20060102150405.000") + "-" + fmt.Sprintf("%d", time.Now().UnixNano())
}