gochat/backend/internal/repository/captain_assistant_response_repo.go

package repository

import (
	"context"
	"strconv"
	"strings"

	"github.com/gochat/gochat/internal/model"
	"github.com/pgvector/pgvector-go"
	"gorm.io/gorm"
)

// CaptainAssistantResponseRepo provides data access for CaptainAssistantResponse.
type CaptainAssistantResponseRepo struct {
	db *gorm.DB
}

func (r *CaptainAssistantResponseRepo) DB() *gorm.DB {
	return r.db
}

func NewCaptainAssistantResponseRepo(db *gorm.DB) *CaptainAssistantResponseRepo {
	return &CaptainAssistantResponseRepo{db: db}
}

func (r *CaptainAssistantResponseRepo) Create(ctx context.Context, resp *model.CaptainAssistantResponse) error {
	return r.db.WithContext(ctx).Create(resp).Error
}

func (r *CaptainAssistantResponseRepo) GetByID(ctx context.Context, id uint) (*model.CaptainAssistantResponse, error) {
	var resp model.CaptainAssistantResponse
	if err := r.db.WithContext(ctx).Preload("Assistant").First(&resp, id).Error; err != nil {
		return nil, err
	}
	return &resp, nil
}

func (r *CaptainAssistantResponseRepo) GetByAccountAndID(ctx context.Context, accountID, id uint) (*model.CaptainAssistantResponse, error) {
	var resp model.CaptainAssistantResponse
	if err := r.db.WithContext(ctx).Preload("Assistant").Where("account_id = ? AND id = ?", accountID, id).First(&resp).Error; err != nil {
		return nil, err
	}
	return &resp, nil
}

func (r *CaptainAssistantResponseRepo) Update(ctx context.Context, resp *model.CaptainAssistantResponse) error {
	return r.db.WithContext(ctx).Save(resp).Error
}

// UpdateEmbedding updates only the embedding column using raw SQL with explicit
// ::vector cast to avoid GORM/pgvector serialization issue (SQLSTATE 42804).
func (r *CaptainAssistantResponseRepo) UpdateEmbedding(ctx context.Context, id uint, embedding pgvector.Vector) error {
	// Format as PG vector literal: [0.1,0.2,...]
	strs := make([]string, len(embedding))
	for i, v := range embedding {
		strs[i] = strconv.FormatFloat(float64(v), 'f', -1, 32)
	}
	vecStr := "[" + strings.Join(strs, ",") + "]"
	return r.db.WithContext(ctx).Exec(
		"UPDATE captain_assistant_responses SET embedding = ?::vector WHERE id = ?",
		vecStr, id,
	).Error
}

func (r *CaptainAssistantResponseRepo) Delete(ctx context.Context, id uint) error {
	return r.db.WithContext(ctx).Delete(&model.CaptainAssistantResponse{}, id).Error
}

func (r *CaptainAssistantResponseRepo) DeleteByAccount(ctx context.Context, accountID, id uint) error {
	return r.db.WithContext(ctx).Where("account_id = ? AND id = ?", accountID, id).Delete(&model.CaptainAssistantResponse{}).Error
}

func (r *CaptainAssistantResponseRepo) ListByAssistant(ctx context.Context, assistantID uint, offset, limit int) ([]model.CaptainAssistantResponse, int64, error) {
	var responses []model.CaptainAssistantResponse
	var count int64
	db := r.db.WithContext(ctx).Model(&model.CaptainAssistantResponse{}).Where("assistant_id = ?", assistantID)
	db.Count(&count)
	if err := db.Offset(offset).Limit(limit).Find(&responses).Error; err != nil {
		return nil, 0, err
	}
	return responses, count, nil
}

func (r *CaptainAssistantResponseRepo) ListByDocument(ctx context.Context, documentableID uint, documentableType string, offset, limit int) ([]model.CaptainAssistantResponse, int64, error) {
	var responses []model.CaptainAssistantResponse
	var count int64
	db := r.db.WithContext(ctx).Model(&model.CaptainAssistantResponse{}).
		Where("documentable_id = ? AND documentable_type = ?", documentableID, documentableType)
	db.Count(&count)
	if err := db.Offset(offset).Limit(limit).Find(&responses).Error; err != nil {
		return nil, 0, err
	}
	return responses, count, nil
}

// SimilaritySearch performs pgvector cosine similarity search.
// Reference: Chatwoot Captain::AssistantResponsesSearchService
func (r *CaptainAssistantResponseRepo) SimilaritySearch(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error) {
	var responses []model.CaptainAssistantResponse
	// Format embedding as PG vector literal for the <=> operator (stub has no driver.Valuer).
	strs := make([]string, len(embedding))
	for i, v := range embedding {
		strs[i] = strconv.FormatFloat(float64(v), 'f', -1, 32)
	}
	vecStr := "[" + strings.Join(strs, ",") + "]"
	// Cosine distance (<=>) orders by closest first.
	// Omit embedding column from SELECT — the pgvector stub can't scan it back.
	if err := r.db.WithContext(ctx).
		Select("captain_assistant_responses.id, captain_assistant_responses.account_id, captain_assistant_responses.assistant_id, captain_assistant_responses.documentable_id, captain_assistant_responses.documentable_type, captain_assistant_responses.question, captain_assistant_responses.answer, captain_assistant_responses.status, captain_assistant_responses.edited, captain_assistant_responses.created_at, captain_assistant_responses.updated_at, COALESCE(captain_documents.external_link, '') AS source_url").
		Joins("JOIN captain_assistants ON captain_assistants.id = captain_assistant_responses.assistant_id AND captain_assistants.account_id = captain_assistant_responses.account_id AND captain_assistants.deleted_at IS NULL").
		Joins("LEFT JOIN captain_documents ON captain_documents.id = captain_assistant_responses.documentable_id AND captain_assistant_responses.documentable_type IN ? AND captain_documents.account_id = captain_assistant_responses.account_id AND captain_documents.assistant_id = captain_assistant_responses.assistant_id AND captain_documents.deleted_at IS NULL", []string{"Captain::Document", "CaptainDocument"}).
		Where("captain_assistant_responses.assistant_id = ? AND captain_assistant_responses.status = ? AND captain_assistant_responses.embedding IS NOT NULL", assistantID, model.ResponseStatusApproved).
		Order(gorm.Expr("captain_assistant_responses.embedding <=> ?::vector", vecStr)).
		Limit(limit).
		Find(&responses).Error; err != nil {
		return nil, err
	}
	return responses, nil
}

// SearchByEmbedding is an alias for SimilaritySearch for backward compatibility.
func (r *CaptainAssistantResponseRepo) SearchByEmbedding(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error) {
	return r.SimilaritySearch(ctx, assistantID, embedding, limit)
}