wraith/internal/ai/conversation.go

package ai

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"time"

	"github.com/google/uuid"
)

// ConversationManager handles CRUD operations for AI conversations stored in SQLite.
type ConversationManager struct {
	db *sql.DB
}

// NewConversationManager creates a manager backed by the given database.
func NewConversationManager(db *sql.DB) *ConversationManager {
	return &ConversationManager{db: db}
}

// Create starts a new conversation and returns its summary.
func (m *ConversationManager) Create(model string) (*ConversationSummary, error) {
	id := uuid.NewString()
	now := time.Now()

	_, err := m.db.Exec(
		`INSERT INTO conversations (id, title, model, messages, tokens_in, tokens_out, created_at, updated_at)
		 VALUES (?, ?, ?, '[]', 0, 0, ?, ?)`,
		id, "New conversation", model, now, now,
	)
	if err != nil {
		return nil, fmt.Errorf("create conversation: %w", err)
	}

	return &ConversationSummary{
		ID:        id,
		Title:     "New conversation",
		Model:     model,
		CreatedAt: now,
		TokensIn:  0,
		TokensOut: 0,
	}, nil
}

// AddMessage appends a message to the conversation's message list.
func (m *ConversationManager) AddMessage(convId string, msg Message) error {
	// Get existing messages
	messages, err := m.GetMessages(convId)
	if err != nil {
		return err
	}

	messages = append(messages, msg)
	data, err := json.Marshal(messages)
	if err != nil {
		return fmt.Errorf("marshal messages: %w", err)
	}

	_, err = m.db.Exec(
		"UPDATE conversations SET messages = ?, updated_at = ? WHERE id = ?",
		string(data), time.Now(), convId,
	)
	if err != nil {
		return fmt.Errorf("update messages: %w", err)
	}

	// Auto-title from first user message
	if len(messages) == 1 && msg.Role == "user" {
		title := extractTitle(msg)
		if title != "" {
			m.db.Exec("UPDATE conversations SET title = ? WHERE id = ?", title, convId)
		}
	}

	return nil
}

// GetMessages returns all messages in a conversation.
func (m *ConversationManager) GetMessages(convId string) ([]Message, error) {
	var messagesJSON string
	err := m.db.QueryRow("SELECT messages FROM conversations WHERE id = ?", convId).Scan(&messagesJSON)
	if err == sql.ErrNoRows {
		return nil, fmt.Errorf("conversation %s not found", convId)
	}
	if err != nil {
		return nil, fmt.Errorf("get messages: %w", err)
	}

	var messages []Message
	if err := json.Unmarshal([]byte(messagesJSON), &messages); err != nil {
		return nil, fmt.Errorf("unmarshal messages: %w", err)
	}
	return messages, nil
}

// List returns all conversations ordered by most recent.
func (m *ConversationManager) List() ([]ConversationSummary, error) {
	rows, err := m.db.Query(
		`SELECT id, title, model, tokens_in, tokens_out, created_at
		 FROM conversations ORDER BY updated_at DESC`,
	)
	if err != nil {
		return nil, fmt.Errorf("list conversations: %w", err)
	}
	defer rows.Close()

	var summaries []ConversationSummary
	for rows.Next() {
		var s ConversationSummary
		var title sql.NullString
		if err := rows.Scan(&s.ID, &title, &s.Model, &s.TokensIn, &s.TokensOut, &s.CreatedAt); err != nil {
			return nil, fmt.Errorf("scan conversation: %w", err)
		}
		if title.Valid {
			s.Title = title.String
		}
		summaries = append(summaries, s)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate conversations: %w", err)
	}

	if summaries == nil {
		summaries = []ConversationSummary{}
	}
	return summaries, nil
}

// Delete removes a conversation and all its messages.
func (m *ConversationManager) Delete(convId string) error {
	result, err := m.db.Exec("DELETE FROM conversations WHERE id = ?", convId)
	if err != nil {
		return fmt.Errorf("delete conversation: %w", err)
	}
	affected, _ := result.RowsAffected()
	if affected == 0 {
		return fmt.Errorf("conversation %s not found", convId)
	}
	return nil
}

// UpdateTokenUsage adds to the token counters for a conversation.
func (m *ConversationManager) UpdateTokenUsage(convId string, tokensIn, tokensOut int) error {
	_, err := m.db.Exec(
		`UPDATE conversations
		 SET tokens_in = tokens_in + ?, tokens_out = tokens_out + ?, updated_at = ?
		 WHERE id = ?`,
		tokensIn, tokensOut, time.Now(), convId,
	)
	if err != nil {
		return fmt.Errorf("update token usage: %w", err)
	}
	return nil
}

// extractTitle generates a title from the first user message (truncated to 80 chars).
func extractTitle(msg Message) string {
	for _, block := range msg.Content {
		if block.Type == "text" && block.Text != "" {
			title := block.Text
			if len(title) > 80 {
				title = title[:77] + "..."
			}
			return title
		}
	}
	return ""
}