chatpb/helpers.go

package chatpb

import (
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/google/uuid"
	"go.wit.com/log"
	"google.golang.org/protobuf/proto"
	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
)

func (c *Chats) AddGeminiComment(s string) *Chat {
	chat := new(Chat)

	chat.From = Who_GEMINI
	chat.Content = s
	chat.Ctime = timestamppb.New(time.Now())

	c.AppendNew(chat)

	return chat
}

func (c *Chats) AddUserComment(s string) *Chat {
	chat := new(Chat)

	chat.From = Who_USER
	chat.Content = s

	c.AppendNew(chat)

	return chat
}

func UnmarshalChats(data []byte) (*Chats, error) {
	c := new(Chats)
	err := c.Unmarshal(data)
	return c, err
}

func UnmarshalChatsTEXT(data []byte) (*Chats, error) {
	c := new(Chats)
	err := c.UnmarshalTEXT(data)
	return c, err
}

func (all *Chats) AddFile(filename string) error {
	// Nil checks for safety.
	if all == nil {
		return fmt.Errorf("cannot call AddFile on a nil *Chats object")
	}
	if all.Chats == nil {
		all.Chats = make([]*Chat, 0)
	}

	data, err := os.ReadFile(filename)
	if err != nil {
		log.Fatalf("Error reading file %s: %v", filename, err)
		return err
	}

	logData, err := UnmarshalChatsTEXT(data)
	if err != nil {
		log.Fatalf("Error unmarshaling log file %s: %v", filename, err)
		return err
	}

	// New logic to handle both flat and nested formats
	for _, chatOrGroup := range logData.GetChats() {
		if len(chatOrGroup.GetEntries()) > 0 {
			// This is a new-style Chat group with entries
			for _, entry := range chatOrGroup.GetEntries() {
				// Convert the ChatEntry into a new Chat object for the flat list.
				newChat := convertEntryToChat(entry, filename)
				if newChat != nil {
					newChat.VerifyUuid()
					all.AppendByUuid(newChat)
				}
			}
		} else {
			// This is an old-style flat Chat entry.
			// We still process it to handle its external content file correctly.
			newChat := convertChatToChat(chatOrGroup, filename)
			if newChat != nil {
				newChat.VerifyUuid()
				all.AppendByUuid(newChat)
			}
		}
	}
	return nil
}

// convertChatToChat handles an old-style Chat message. It creates a clean
// copy and resolves its external content file.
func convertChatToChat(chat *Chat, filename string) *Chat {
	if chat == nil {
		return nil
	}

	// Manually create a ChatEntry from the Chat fields to reuse the logic.
	entry := &ChatEntry{
		From:        chat.GetFrom(),
		Ctime:       chat.GetCtime(),
		Content:     chat.GetContent(),
		Table:       chat.GetTable(),
		ToolCalls:   chat.GetToolCalls(),
		ContentFile: chat.GetContentFile(),
		Uuid:        chat.GetUuid(),
		Snippets:    chat.GetSnippets(),
	}
	return convertEntryToChat(entry, filename)
}

// convertEntryToChat creates a new Chat object from a ChatEntry's data
// and resolves its external content file.
func convertEntryToChat(entry *ChatEntry, filename string) *Chat {
	if entry == nil {
		return nil
	}

	// Create a new Chat object and copy the fields.
	newChat := &Chat{
		From:      entry.GetFrom(),
		Ctime:     entry.GetCtime(),
		Table:     entry.GetTable(),
		ToolCalls: entry.GetToolCalls(),
		Snippets:  entry.GetSnippets(),
		Uuid:      entry.GetUuid(),
	}

	// Handle content: prefer content_file, fallback to content.
	var content string
	if contentFile := entry.GetContentFile(); contentFile != "" {
		logDir := filepath.Dir(filename)
		contentPath := filepath.Join(logDir, contentFile)
		contentBytes, err := os.ReadFile(contentPath)
		if err != nil {
			content = fmt.Sprintf("--- ERROR: Could not read content file %s: %v ---", contentPath, err)
		} else {
			content = string(contentBytes)
		}
	} else {
		content = entry.GetContent()
	}
	newChat.Content = content

	return newChat
}

func (chats *Chats) VerifyUuids() bool {
	var changed bool

	all := chats.SortByUuid()
	for all.Scan() {
		chat := all.Next()
		if chat.Uuid == "" {
			chat.Uuid = uuid.New().String()
			changed = true
		}
	}
	return changed
}

func (c *Chat) VerifyUuid() bool {
	if c.Uuid == "" {
		c.Uuid = uuid.New().String()
		return true
	}
	return false
}

func (x *Chats) AppendNew(y *Chat) {
	x.Lock()
	defer x.Unlock()

	var chat *Chat
	chat = proto.Clone(y).(*Chat)

	x.Chats = append(x.Chats, chat)
}