package repository

import (
	"context"

	"git.techease.ru/Smart-search/smart-search-back/internal/model"
	errs "git.techease.ru/Smart-search/smart-search-back/pkg/errors"
	sq "github.com/Masterminds/squirrel"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type tokenUsageRepository struct {
	pool *pgxpool.Pool
	qb   sq.StatementBuilderType
}

func NewTokenUsageRepository(pool *pgxpool.Pool) TokenUsageRepository {
	return &tokenUsageRepository{
		pool: pool,
		qb:   sq.StatementBuilder.PlaceholderFormat(sq.Dollar),
	}
}

func (r *tokenUsageRepository) Create(ctx context.Context, usage *model.TokenUsage) error {
	return r.createWithExecutor(ctx, r.pool, usage)
}

func (r *tokenUsageRepository) CreateTx(ctx context.Context, tx pgx.Tx, usage *model.TokenUsage) error {
	return r.createWithExecutor(ctx, tx, usage)
}

func (r *tokenUsageRepository) createWithExecutor(ctx context.Context, exec DBTX, usage *model.TokenUsage) error {
	query := r.qb.Insert("request_token_usage").Columns(
		"request_id", "request_token_count", "response_token_count", "token_cost", "type",
	).Values(
		usage.RequestID, usage.RequestTokenCount, usage.ResponseTokenCount, usage.TokenCost, usage.Type,
	).Suffix("RETURNING id, created_at")

	sqlQuery, args, err := query.ToSql()
	if err != nil {
		return errs.NewInternalError(errs.DatabaseError, "failed to build query", err)
	}

	err = exec.QueryRow(ctx, sqlQuery, args...).Scan(&usage.ID, &usage.CreatedAt)
	if err != nil {
		return errs.NewInternalError(errs.DatabaseError, "failed to create token usage", err)
	}

	return nil
}