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 maxTokensApprox = 100000
const summarizeThreshold = 80000
const charsPerToken = 4

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

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
	realTokens int
}

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) Clear() {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	cs.conv.Messages = []FeedMessage{}
	cs.conv.Summary = ""
	cs.conv.CreatedAt = time.Now().Format(time.RFC3339)
	cs.conv.UpdatedAt = time.Now().Format(time.RFC3339)
	cs.realTokens = 0
	cs.save()
}

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

func (cs *ConversationStore) TrimOld(keepCount int) {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	if len(cs.conv.Messages) <= keepCount {
		return
	}
	cs.conv.Messages = cs.conv.Messages[len(cs.conv.Messages)-keepCount:]
	cs.save()
}

func (cs *ConversationStore) ApproxTokenCount() int {
	if cs.realTokens > 0 {
		return cs.realTokens
	}
	return cs.ApproxTokenCountDetailed().total
}

// AddRealTokens accumulates actual token counts from the API response.
func (cs *ConversationStore) AddRealTokens(tokens int) {
	if tokens <= 0 {
		return
	}
	cs.mu.Lock()
	cs.realTokens += tokens
	cs.mu.Unlock()
}

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 {
		count := utf8.RuneCountInString(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() > summarizeThreshold
}

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())
}