smart-search-back/internal/service/request.go

package service

import (
	"context"
	"math"

	"git.techease.ru/Smart-search/smart-search-back/internal/ai"
	"git.techease.ru/Smart-search/smart-search-back/internal/model"
	"git.techease.ru/Smart-search/smart-search-back/internal/repository"
	"git.techease.ru/Smart-search/smart-search-back/pkg/errors"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

type requestService struct {
	requestRepo    repository.RequestRepository
	supplierRepo   repository.SupplierRepository
	tokenUsageRepo repository.TokenUsageRepository
	userRepo       repository.UserRepository
	openAI         *ai.OpenAIClient
	perplexity     *ai.PerplexityClient
	txManager      *repository.TxManager
}

func NewRequestService(
	requestRepo repository.RequestRepository,
	supplierRepo repository.SupplierRepository,
	tokenUsageRepo repository.TokenUsageRepository,
	userRepo repository.UserRepository,
	openAI *ai.OpenAIClient,
	perplexity *ai.PerplexityClient,
	txManager *repository.TxManager,
) RequestService {
	return &requestService{
		requestRepo:    requestRepo,
		supplierRepo:   supplierRepo,
		tokenUsageRepo: tokenUsageRepo,
		userRepo:       userRepo,
		openAI:         openAI,
		perplexity:     perplexity,
		txManager:      txManager,
	}
}

func (s *requestService) CreateTZ(ctx context.Context, userID int, requestTxt string) (uuid.UUID, string, error) {
	req := &model.Request{
		UserID:     userID,
		RequestTxt: requestTxt,
	}

	if err := s.requestRepo.Create(ctx, req); err != nil {
		return uuid.Nil, "", err
	}

	if requestTxt == "" {
		return req.ID, "", nil
	}

	tzText, err := s.openAI.GenerateTZ(requestTxt)
	if err != nil {
		if err := s.requestRepo.UpdateWithTZ(ctx, req.ID, "", false); err != nil {
			return req.ID, "", err
		}
		return req.ID, "", err
	}

	inputLen := len(requestTxt)
	outputLen := len(tzText)
	promptTokens := 500

	inputTokens := int(math.Ceil(float64(inputLen) / 2.0))
	outputTokens := int(math.Ceil(float64(outputLen) / 2.0))

	totalTokens := inputTokens + outputTokens + promptTokens
	tokenPrice := 25000.0 / 1000000.0
	cost := float64(totalTokens) * tokenPrice

	tokenUsage := &model.TokenUsage{
		RequestID:          req.ID,
		RequestTokenCount:  inputTokens + promptTokens,
		ResponseTokenCount: outputTokens,
		TokenCost:          cost,
		Type:               "tz",
	}

	err = s.txManager.WithTx(ctx, func(tx pgx.Tx) error {
		if err := s.tokenUsageRepo.CreateTx(ctx, tx, tokenUsage); err != nil {
			return err
		}

		if err := s.userRepo.UpdateBalanceTx(ctx, tx, userID, -cost); err != nil {
			return err
		}

		if err := s.requestRepo.UpdateWithTZTx(ctx, tx, req.ID, tzText, true); err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return req.ID, "", err
	}

	return req.ID, tzText, nil
}

func (s *requestService) ApproveTZ(ctx context.Context, requestID uuid.UUID, tzText string, userID int) ([]*model.Supplier, error) {
	isOwner, err := s.requestRepo.CheckOwnership(ctx, requestID, userID)
	if err != nil {
		return nil, err
	}
	if !isOwner {
		return nil, errors.NewBusinessError(errors.PermissionDenied, "access denied to this request")
	}

	if err = s.requestRepo.UpdateFinalTZ(ctx, requestID, tzText); err != nil {
		return nil, err
	}

	var suppliers []*model.Supplier
	var promptTokens, responseTokens int

	for attempt := 0; attempt < 3; attempt++ {
		suppliers, promptTokens, responseTokens, err = s.perplexity.FindSuppliers(tzText)
		if err == nil {
			break
		}
	}

	if err != nil {
		return nil, err
	}

	if len(suppliers) == 0 {
		return nil, errors.NewInternalError(errors.AIAPIError, "no suppliers found", nil)
	}

	tokenPrice := 25000.0 / 1000000.0
	totalTokens := promptTokens + responseTokens
	cost := float64(totalTokens) * tokenPrice

	tokenUsage := &model.TokenUsage{
		RequestID:          requestID,
		RequestTokenCount:  promptTokens,
		ResponseTokenCount: responseTokens,
		TokenCost:          cost,
		Type:               "suppliers",
	}

	err = s.txManager.WithTx(ctx, func(tx pgx.Tx) error {
		if err := s.supplierRepo.BulkInsertTx(ctx, tx, requestID, suppliers); err != nil {
			return err
		}

		if err := s.tokenUsageRepo.CreateTx(ctx, tx, tokenUsage); err != nil {
			return err
		}

		if err := s.userRepo.UpdateBalanceTx(ctx, tx, userID, -cost); err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return suppliers, nil
}

func (s *requestService) GetMailingList(ctx context.Context, userID int) ([]*model.Request, error) {
	return s.requestRepo.GetByUserID(ctx, userID)
}

func (s *requestService) GetMailingListByID(ctx context.Context, requestID uuid.UUID, userID int) (*model.RequestDetail, error) {
	isOwner, err := s.requestRepo.CheckOwnership(ctx, requestID, userID)
	if err != nil {
		return nil, err
	}
	if !isOwner {
		return nil, errors.NewBusinessError(errors.PermissionDenied, "access denied to this request")
	}

	return s.requestRepo.GetDetailByID(ctx, requestID)
}