core/server/store/postgres/message.go

package postgres

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

	"chat-switchboard/models"
	"chat-switchboard/store"
)

type MessageStore struct{}

func NewMessageStore() *MessageStore { return &MessageStore{} }

func (s *MessageStore) Create(ctx context.Context, m *models.Message) error {
	toolCallsJSON := ToJSON(m.ToolCalls)
	metadataJSON := ToJSON(m.Metadata)
	return DB.QueryRowContext(ctx, `
		INSERT INTO messages (channel_id, role, content, model, tokens_used, tool_calls,
			metadata, parent_id, sibling_index, participant_type, participant_id)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)
		RETURNING id, created_at`,
		m.ChannelID, m.Role, m.Content, m.Model, m.TokensUsed,
		toolCallsJSON, metadataJSON,
		models.NullString(m.ParentID), m.SiblingIndex,
		m.ParticipantType, m.ParticipantID,
	).Scan(&m.ID, &m.CreatedAt)
}

func (s *MessageStore) GetByID(ctx context.Context, id string) (*models.Message, error) {
	var m models.Message
	var parentID sql.NullString
	var toolCallsJSON, metadataJSON []byte
	var deletedAt sql.NullTime
	err := DB.QueryRowContext(ctx, `
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM messages WHERE id = $1`, id).Scan(
		&m.ID, &m.ChannelID, &m.Role, &m.Content, &m.Model, &m.TokensUsed,
		&toolCallsJSON, &metadataJSON,
		&parentID, &m.SiblingIndex, &m.ParticipantType, &m.ParticipantID,
		&deletedAt, &m.CreatedAt,
	)
	if err != nil {
		return nil, err
	}
	m.ParentID = NullableStringPtr(parentID)
	json.Unmarshal(toolCallsJSON, &m.ToolCalls)
	json.Unmarshal(metadataJSON, &m.Metadata)
	if deletedAt.Valid {
		m.DeletedAt = &deletedAt.Time
	}
	return &m, nil
}

func (s *MessageStore) Update(ctx context.Context, id string, fields map[string]interface{}) error {
	b := NewUpdate("messages")
	for k, v := range fields {
		if k == "tool_calls" || k == "metadata" {
			b.SetJSON(k, v)
		} else {
			b.Set(k, v)
		}
	}
	if !b.HasSets() {
		return nil
	}
	b.Where("id", id)
	_, err := b.Exec(DB)
	return err
}

func (s *MessageStore) Delete(ctx context.Context, id string) error {
	now := time.Now()
	_, err := DB.ExecContext(ctx, "UPDATE messages SET deleted_at = $1 WHERE id = $2", now, id)
	return err
}

func (s *MessageStore) ListForChannel(ctx context.Context, channelID string, opts store.ListOptions) ([]models.Message, error) {
	b := NewSelect(
		"id, channel_id, role, content, model, tokens_used, tool_calls, metadata, parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at",
		"messages",
	).Where("channel_id = ?", channelID).WhereRaw("deleted_at IS NULL")
	if opts.Sort == "" {
		b.OrderBy("created_at", "ASC")
	}
	b.Paginate(opts)

	q, args := b.Build()
	rows, err := DB.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetChildren(ctx context.Context, parentID string) ([]models.Message, error) {
	rows, err := DB.QueryContext(ctx, `
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM messages WHERE parent_id = $1 AND deleted_at IS NULL
		ORDER BY sibling_index`, parentID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetSiblings(ctx context.Context, messageID string) ([]models.Message, error) {
	rows, err := DB.QueryContext(ctx, `
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM messages
		WHERE parent_id = (SELECT parent_id FROM messages WHERE id = $1)
			AND deleted_at IS NULL
		ORDER BY sibling_index`, messageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetPathToRoot(ctx context.Context, messageID string) ([]models.Message, error) {
	rows, err := DB.QueryContext(ctx, `
		WITH RECURSIVE path AS (
			SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
				parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
			FROM messages WHERE id = $1
			UNION ALL
			SELECT m.id, m.channel_id, m.role, m.content, m.model, m.tokens_used, m.tool_calls, m.metadata,
				m.parent_id, m.sibling_index, m.participant_type, m.participant_id, m.deleted_at, m.created_at
			FROM messages m JOIN path p ON m.id = p.parent_id
		)
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM path ORDER BY created_at ASC`, messageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetNextSiblingIndex(ctx context.Context, parentID string) (int, error) {
	var maxIdx sql.NullInt64
	err := DB.QueryRowContext(ctx,
		"SELECT MAX(sibling_index) FROM messages WHERE parent_id = $1 AND deleted_at IS NULL",
		parentID).Scan(&maxIdx)
	if err != nil || !maxIdx.Valid {
		return 0, err
	}
	return int(maxIdx.Int64) + 1, nil
}

func (s *MessageStore) CountForChannel(ctx context.Context, channelID string) (int, error) {
	var count int
	err := DB.QueryRowContext(ctx,
		"SELECT COUNT(*) FROM messages WHERE channel_id = $1 AND deleted_at IS NULL",
		channelID).Scan(&count)
	return count, err
}

func scanMessages(rows *sql.Rows) ([]models.Message, error) {
	var result []models.Message
	for rows.Next() {
		var m models.Message
		var parentID sql.NullString
		var toolCallsJSON, metadataJSON []byte
		var deletedAt sql.NullTime
		err := rows.Scan(
			&m.ID, &m.ChannelID, &m.Role, &m.Content, &m.Model, &m.TokensUsed,
			&toolCallsJSON, &metadataJSON,
			&parentID, &m.SiblingIndex, &m.ParticipantType, &m.ParticipantID,
			&deletedAt, &m.CreatedAt,
		)
		if err != nil {
			return nil, err
		}
		m.ParentID = NullableStringPtr(parentID)
		json.Unmarshal(toolCallsJSON, &m.ToolCalls)
		json.Unmarshal(metadataJSON, &m.Metadata)
		if deletedAt.Valid {
			m.DeletedAt = &deletedAt.Time
		}
		result = append(result, m)
	}
	return result, rows.Err()
}

// ── CS2 additions (v0.29.0) ─────────────────────────────────────────────

func (s *MessageStore) CountAll(ctx context.Context) (int, error) {
	var count int
	err := DB.QueryRowContext(ctx, "SELECT COUNT(*) FROM messages").Scan(&count)
	return count, err
}

// ── CS5c additions (v0.29.0) ─────────────────────────────────────────────

func (s *MessageStore) SearchInChannel(ctx context.Context, channelID, query, roleFilter string, limit int) ([]store.ChannelSearchResult, error) {
	roleClause := ""
	queryArgs := []interface{}{channelID, query, limit}
	if roleFilter == "user" || roleFilter == "assistant" {
		roleClause = "AND m.role = $4"
		queryArgs = append(queryArgs, roleFilter)
	}

	rows, err := DB.QueryContext(ctx, fmt.Sprintf(`
		SELECT m.id, m.role,
		       ts_headline('english', m.content, plainto_tsquery('english', $2),
		                   'MaxWords=60, MinWords=20, StartSel=**, StopSel=**') AS headline,
		       ts_rank(to_tsvector('english', m.content), plainto_tsquery('english', $2)) AS rank,
		       m.created_at
		FROM messages m
		WHERE m.channel_id = $1
		  AND m.deleted_at IS NULL
		  AND m.role IN ('user', 'assistant')
		  AND to_tsvector('english', m.content) @@ plainto_tsquery('english', $2)
		  %s
		ORDER BY rank DESC, m.created_at DESC
		LIMIT $3
	`, roleClause), queryArgs...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	results := make([]store.ChannelSearchResult, 0)
	for rows.Next() {
		var r store.ChannelSearchResult
		if err := rows.Scan(&r.MessageID, &r.Role, &r.Excerpt, &r.Rank, &r.Timestamp); err != nil {
			continue
		}
		results = append(results, r)
	}
	return results, rows.Err()
}

// ── CS7a additions (v0.29.0) ────────────────────────────────────────────

func (s *MessageStore) ListWithSenderInfo(ctx context.Context, channelID string, limit, offset int) ([]store.MessageWithSender, int, error) {
	var total int
	err := DB.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM messages WHERE channel_id = $1 AND deleted_at IS NULL`,
		channelID).Scan(&total)
	if err != nil {
		return nil, 0, err
	}

	rows, err := DB.QueryContext(ctx, `
		SELECT m.id, m.channel_id, m.role, m.content, m.model, m.tokens_used, m.parent_id,
		       m.sibling_index, m.participant_type, m.participant_id,
		       CASE WHEN m.participant_type = 'user' THEN COALESCE(NULLIF(u.display_name, ''), u.username)
		            WHEN m.participant_type = 'persona' THEN p.name
		            ELSE NULL END AS sender_name,
		       CASE WHEN m.participant_type = 'user' THEN u.avatar_url
		            WHEN m.participant_type = 'persona' THEN p.avatar
		            ELSE NULL END AS sender_avatar,
		       m.created_at
		FROM messages m
		LEFT JOIN users u ON m.participant_type = 'user' AND m.participant_id = u.id::text
		LEFT JOIN personas p ON m.participant_type = 'persona' AND m.participant_id = p.id::text
		WHERE m.channel_id = $1 AND m.deleted_at IS NULL
		ORDER BY m.created_at ASC
		LIMIT $2 OFFSET $3
	`, channelID, limit, offset)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	results := make([]store.MessageWithSender, 0)
	for rows.Next() {
		var m store.MessageWithSender
		if err := rows.Scan(
			&m.ID, &m.ChannelID, &m.Role, &m.Content,
			&m.Model, &m.TokensUsed, &m.ParentID,
			&m.SiblingIndex, &m.ParticipantType, &m.ParticipantID,
			&m.SenderName, &m.SenderAvatar,
			&m.CreatedAt,
		); err != nil {
			return nil, total, err
		}
		results = append(results, m)
	}
	return results, total, rows.Err()
}

func (s *MessageStore) GetParentAndRole(ctx context.Context, messageID, channelID string) (*string, string, error) {
	var parentID sql.NullString
	var role string
	err := DB.QueryRowContext(ctx, `
		SELECT parent_id, role FROM messages
		WHERE id = $1 AND channel_id = $2 AND deleted_at IS NULL
	`, messageID, channelID).Scan(&parentID, &role)
	if err != nil {
		return nil, "", err
	}
	return NullableStringPtr(parentID), role, nil
}