package query

import (
	"context"
	"time"

	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Operations"
	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
	"google.golang.org/grpc"

	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/gtrace"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/safe"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/types"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
	"github.com/ydb-platform/ydb-go-sdk/v3/query"
	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
)

//go:generate mockgen -destination grpc_client_mock_test.go --typed -package query -write_package_comment=false github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1 QueryServiceClient,QueryService_AttachSessionClient,QueryService_ExecuteQueryClient

var (
	_ query.Client = (*Client)(nil)
	_ sessionPool  = (*pool.Pool[*Session, Session])(nil)
)

type (
	sessionPool interface {
		closer.Closer

		Stats() pool.Stats
		With(ctx context.Context, f func(ctx context.Context, s *Session) error, opts ...retry.Option) error
	}
	closeState struct {
		closed       bool
		nextCancelID uint64
		cancels      map[uint64]context.CancelFunc
	}
	Client struct {
		config              *config.Config
		client              Ydb_Query_V1.QueryServiceClient
		explicitSessionPool sessionPool

		// implicitSessionPool is a pool of implicit sessions,
		// i.e. fake sessions created without CreateSession/AttachSession requests.
		implicitSessionPool sessionPool

		closed *xsync.Value[*closeState]
	}
)

func (c *Client) Config() *config.Config {
	return c.config
}

func fetchScriptResults(ctx context.Context,
	client Ydb_Query_V1.QueryServiceClient,
	opID string, opts ...options.FetchScriptOption,
) (*options.FetchScriptResult, error) {
	r, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
		request := &options.FetchScriptResultsRequest{
			FetchScriptResultsRequest: Ydb_Query.FetchScriptResultsRequest{
				OperationId: opID,
			},
		}
		for _, opt := range opts {
			if opt != nil {
				opt(request)
			}
		}

		response, err := client.FetchScriptResults(ctx, &request.FetchScriptResultsRequest)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		rs := response.GetResultSet()
		columns := rs.GetColumns()
		columnNames := make([]string, len(columns))
		columnTypes := make([]types.Type, len(columns))
		for i := range columns {
			columnNames[i] = columns[i].GetName()
			columnTypes[i] = types.TypeFromYDB(columns[i].GetType())
		}
		rows := make([]query.Row, len(rs.GetRows()))
		for i, r := range rs.GetRows() {
			rows[i] = NewRow(columns, r)
		}

		return &options.FetchScriptResult{
			ResultSetIndex: response.GetResultSetIndex(),
			ResultSet:      MaterializedResultSet(int(response.GetResultSetIndex()), columnNames, columnTypes, rows),
			NextToken:      response.GetNextFetchToken(),
		}, nil
	}, retry.WithIdempotent(true))
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

func (c *Client) FetchScriptResults(ctx context.Context,
	opID string, opts ...options.FetchScriptOption,
) (*options.FetchScriptResult, error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return nil, err
	}
	defer leave()

	r, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
		r, err := fetchScriptResults(ctx, c.client, opID,
			append(opts, func(request *options.FetchScriptResultsRequest) {
				request.Trace = c.config.Trace()
			})...,
		)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		return r, nil
	}, retry.WithIdempotent(true))
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

type executeScriptSettings struct {
	executeSettings

	ttl             time.Duration
	operationParams *Ydb_Operations.OperationParams
}

func (s *executeScriptSettings) OperationParams() *Ydb_Operations.OperationParams {
	return s.operationParams
}

func (s *executeScriptSettings) ResultsTTL() time.Duration {
	return s.ttl
}

func executeScript(ctx context.Context,
	client Ydb_Query_V1.QueryServiceClient, request *Ydb_Query.ExecuteScriptRequest, grpcOpts ...grpc.CallOption,
) (*options.ExecuteScriptOperation, error) {
	op, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.ExecuteScriptOperation, error) {
		response, err := client.ExecuteScript(ctx, request, grpcOpts...)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		return &options.ExecuteScriptOperation{
			ID:            response.GetId(),
			ConsumedUnits: response.GetCostInfo().GetConsumedUnits(),
			Metadata:      options.ToMetadataExecuteQuery(response.GetMetadata()),
		}, nil
	}, retry.WithIdempotent(true))
	if err != nil {
		return op, xerrors.WithStackTrace(err)
	}

	return op, nil
}

func (c *Client) ExecuteScript(
	ctx context.Context, q string, ttl time.Duration, opts ...options.Execute,
) (
	op *options.ExecuteScriptOperation, err error,
) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return nil, err
	}
	defer leave()

	settings := &executeScriptSettings{
		executeSettings: options.ExecuteSettings(opts...),
		ttl:             ttl,
		operationParams: operation.Params(
			ctx,
			c.config.OperationTimeout(),
			c.config.OperationCancelAfter(),
			operation.ModeSync,
		),
	}

	request, grpcOpts, err := executeQueryScriptRequest(q, settings)
	if err != nil {
		return op, xerrors.WithStackTrace(err)
	}

	op, err = executeScript(ctx, c.client, request, grpcOpts...)
	if err != nil {
		return op, xerrors.WithStackTrace(err)
	}

	return op, nil
}

func (c *Client) Close(ctx context.Context) error {
	if c == nil {
		return xerrors.WithStackTrace(errNilClient)
	}

	var cancels []context.CancelFunc
	c.closed.Change(func(old *closeState) *closeState {
		if old.closed {
			return old
		}

		old.closed = true
		cancels = make([]context.CancelFunc, 0, len(old.cancels))
		for _, cancel := range old.cancels {
			cancels = append(cancels, cancel)
		}
		old.cancels = nil

		return old
	})
	if cancels == nil {
		return nil
	}

	for _, cancel := range cancels {
		cancel()
	}

	if err := c.explicitSessionPool.Close(ctx); err != nil {
		return xerrors.WithStackTrace(err)
	}

	if err := c.implicitSessionPool.Close(ctx); err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func (c *Client) enter(ctx context.Context) (context.Context, context.CancelFunc, error) {
	if c == nil {
		return nil, nil, xerrors.WithStackTrace(errNilClient)
	}

	ctx, cancel := context.WithCancel(ctx)

	var (
		id     uint64
		closed bool
	)
	c.closed.Change(func(old *closeState) *closeState {
		if old.closed {
			closed = true

			return old
		}

		new := *old

		new.nextCancelID++
		id = new.nextCancelID
		new.cancels[id] = cancel

		return &new
	})
	if closed {
		cancel()

		return nil, nil, xerrors.WithStackTrace(errClosedClient)
	}

	return ctx, func() {
		c.unregisterCloseCancel(id)
		cancel()
	}, nil
}

func (c *Client) registerCloseCancel(cancel context.CancelFunc) func() {
	var id uint64
	c.closed.Change(func(old *closeState) *closeState {
		if old.closed {
			return old
		}

		new := *old

		new.nextCancelID++
		id = new.nextCancelID
		new.cancels[id] = cancel

		return &new
	})
	if id == 0 {
		cancel()

		return func() {}
	}

	return func() {
		c.unregisterCloseCancel(id)
	}
}

func (c *Client) unregisterCloseCancel(id uint64) {
	c.closed.Change(func(old *closeState) *closeState {
		delete(old.cancels, id)

		return old
	})
}

func do(
	ctx context.Context,
	pool sessionPool,
	op func(ctx context.Context, s *Session) error,
	opts ...retry.Option,
) (finalErr error) {
	err := pool.With(ctx, func(ctx context.Context, s *Session) error {
		s.SetStatus(StatusInUse)

		err := op(ctx, s)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		s.SetStatus(StatusIdle)

		return nil
	}, opts...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func (c *Client) Do(ctx context.Context, op query.Operation, opts ...options.DoOption) (finalErr error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}
	defer leave()

	var (
		settings = options.ParseDoOpts(c.config.Trace(), opts...)
		onDone   = gtrace.QueryOnDo(settings.Trace(), &ctx,
			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Do"),
			settings.Label(),
		)
		attempts = 0
	)
	defer func() {
		onDone(attempts, finalErr)
	}()

	err = do(ctx, c.explicitSessionPool,
		func(ctx context.Context, s *Session) error {
			return op(ctx, s)
		},
		append([]retry.Option{
			// Driver-level trace.Retry (e.g. spans.WithTraces) so retry
			// callers see ydb.RunWithRetry / ydb.Try spans for Do().
			retry.WithTrace(c.config.TraceRetry()),
			retry.WithTrace(&trace.Retry{
				OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
					return func(info trace.RetryLoopDoneInfo) {
						attempts = info.Attempts
					}
				},
			}),
		}, settings.RetryOpts()...)...,
	)

	return err
}

func doTx(
	ctx context.Context,
	pool sessionPool,
	op query.TxOperation,
	txSettings tx.Settings,
	opts ...retry.Option,
) (finalErr error) {
	err := do(ctx, pool, func(ctx context.Context, s *Session) (opErr error) {
		tx, err := s.Begin(ctx, txSettings)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		defer func() {
			_ = tx.Rollback(ctx)
		}()

		err = op(ctx, tx)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		err = tx.CommitTx(ctx)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, opts...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

// validateTxControl checks if user provided a TxControl with BeginTx but without CommitTx
func validateTxControl(settings executeSettings) error {
	if settings.UserProvidedTxControl() && settings.TxControl() != nil {
		txControl, ok := settings.TxControl().(*tx.Control)
		if ok && txControl.IsBeginTxWithoutCommit() {
			return xerrors.WithStackTrace(query.ErrTxControlWithoutCommit)
		}
	}

	return nil
}

func clientQueryRow(
	ctx context.Context, pool sessionPool, q string, settings executeSettings, resultOpts ...resultOption,
) (row query.Row, finalErr error) {
	if err := validateTxControl(settings); err != nil {
		return nil, err
	}

	err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
		row, err = s.queryRow(ctx, q, settings, resultOpts...)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return row, nil
}

// QueryRow is a helper which read only one row from first result set in result
func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute) (_ query.Row, finalErr error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}
	defer leave()

	settings := options.ExecuteSettings(opts...)

	onDone := gtrace.QueryOnQueryRow(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).QueryRow"),
		q, settings.Label(),
	)
	defer func() {
		onDone(finalErr)
	}()

	row, err := clientQueryRow(ctx, c.pool(), q, settings,
		withStreamResultTrace(c.config.Trace()),
		withIssuesHandler(settings.IssuesOpts()),
	)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return row, nil
}

func clientExec(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (finalErr error) {
	settings := options.ExecuteSettings(opts...)

	if err := validateTxControl(settings); err != nil {
		return err
	}

	err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
		streamResult, err := s.execute(ctx, q, settings,
			withStreamResultTrace(s.trace), withIssuesHandler(settings.IssuesOpts()))
		if err != nil {
			return xerrors.WithStackTrace(err)
		}
		defer func() {
			_ = streamResult.Close(ctx)
		}()

		err = readAll(ctx, streamResult)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func (c *Client) Exec(ctx context.Context, q string, opts ...options.Execute) (finalErr error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}
	defer leave()

	settings := options.ExecuteSettings(opts...)
	onDone := gtrace.QueryOnExec(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Exec"),
		q,
		settings.Label(),
	)
	defer func() {
		onDone(finalErr)
	}()

	err = clientExec(ctx, c.pool(), q, opts...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (
	r query.Result, err error,
) {
	settings := options.ExecuteSettings(opts...)

	if err := validateTxControl(settings); err != nil {
		return nil, err
	}

	err = do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
		streamResult, err := s.execute(ctx, q, settings,
			withStreamResultTrace(s.trace), withIssuesHandler(settings.IssuesOpts()))
		if err != nil {
			return xerrors.WithStackTrace(err)
		}
		defer func() {
			_ = streamResult.Close(ctx)
		}()

		r, err = resultToMaterializedResult(ctx, streamResult)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

func (c *Client) Query(ctx context.Context, q string, opts ...options.Execute) (r query.Result, err error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}
	defer leave()

	settings := options.ExecuteSettings(opts...)
	onDone := gtrace.QueryOnQuery(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Query"),
		q, settings.Label(),
	)
	defer func() {
		onDone(err)
	}()

	r, err = clientQuery(ctx, c.pool(), q, opts...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

func clientQueryResultSet(
	ctx context.Context, pool sessionPool, q string, settings executeSettings, resultOpts ...resultOption,
) (rs result.ClosableResultSet, rowsCount int, finalErr error) {
	if err := validateTxControl(settings); err != nil {
		return nil, 0, err
	}

	err := do(ctx, pool, func(ctx context.Context, s *Session) error {
		streamResult, err := s.execute(ctx, q, settings, resultOpts...)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}
		defer func() {
			_ = streamResult.Close(ctx)
		}()

		rs, rowsCount, err = readMaterializedResultSet(ctx, streamResult)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return nil, 0, xerrors.WithStackTrace(err)
	}

	return rs, rowsCount, nil
}

// QueryResultSet is a helper which read all rows from first result set in result
func (c *Client) QueryResultSet(
	ctx context.Context, q string, opts ...options.Execute,
) (rs result.ClosableResultSet, finalErr error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}
	defer leave()

	var (
		settings  = options.ExecuteSettings(opts...)
		rowsCount int
	)

	onDone := gtrace.QueryOnQueryResultSet(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).QueryResultSet"),
		q, settings.Label(),
	)
	defer func() {
		onDone(finalErr, rowsCount)
	}()

	rs, rowsCount, err = clientQueryResultSet(ctx, c.pool(), q, settings,
		withStreamResultTrace(c.config.Trace()),
		withIssuesHandler(settings.IssuesOpts()),
	)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return rs, nil
}

// pool returns the appropriate session pool based on the client configuration.
// If implicit sessions are enabled, it returns the implicit session pool;
// otherwise, it returns the explicit session pool.
func (c *Client) pool() sessionPool {
	if c.config.AllowImplicitSessions() {
		return c.implicitSessionPool
	}

	return c.explicitSessionPool
}

func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options.DoTxOption) (finalErr error) {
	ctx, leave, err := c.enter(ctx)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}
	defer leave()

	var (
		settings = options.ParseDoTxOpts(c.config.Trace(), opts...)
		onDone   = gtrace.QueryOnDoTx(settings.Trace(), &ctx,
			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).DoTx"),
			settings.Label(),
		)
		attempts = 0
	)
	defer func() {
		onDone(attempts, finalErr)
	}()

	if lazyTx := settings.LazyTx(); lazyTx != nil {
		ctx = tx.WithLazyTx(ctx, *lazyTx)
	}

	err = doTx(ctx, c.explicitSessionPool, op,
		settings.TxSettings(),
		append(
			[]retry.Option{
				// Driver-level trace.Retry (e.g. spans.WithTraces) so retry
				// callers see ydb.RunWithRetry / ydb.Try spans for DoTx().
				retry.WithTrace(c.config.TraceRetry()),
				retry.WithTrace(&trace.Retry{
					OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
						return func(info trace.RetryLoopDoneInfo) {
							attempts = info.Attempts
						}
					},
				}),
			},
			settings.RetryOpts()...,
		)...,
	)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func CreateSession(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, cfg *config.Config) (*Session, error) {
	s, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*Session, error) {
		var (
			createCtx    context.Context
			cancelCreate context.CancelFunc
		)
		if d := cfg.SessionCreateTimeout(); d > 0 {
			createCtx, cancelCreate = xcontext.WithTimeout(ctx, d)
		} else {
			createCtx, cancelCreate = xcontext.WithCancel(ctx)
		}
		defer cancelCreate()

		s, err := createSession(createCtx, client,
			WithDeleteTimeout(cfg.SessionDeleteTimeout()),
			WithTrace(cfg.Trace()),
		)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		s.lazyTx = cfg.LazyTx()

		return s, nil
	})
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return s, nil
}

func New(ctx context.Context, cc grpc.ClientConnInterface, cfg *config.Config) (*Client, error) {
	onDone := gtrace.QueryOnNew(cfg.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.New"),
	)
	defer onDone()

	client := Ydb_Query_V1.NewQueryServiceClient(cc)

	return newWithQueryServiceClient(ctx, client, cc, cfg)
}

//nolint:funlen
func newWithQueryServiceClient(ctx context.Context,
	client Ydb_Query_V1.QueryServiceClient,
	cc grpc.ClientConnInterface,
	cfg *config.Config,
) (*Client, error) {
	c := &Client{
		config: cfg,
		client: client,
		closed: xsync.NewValue(&closeState{
			cancels: make(map[uint64]context.CancelFunc),
		}),
	}
	explicitSessionPool, err := pool.New(ctx,
		pool.WithLimit[*Session](cfg.PoolLimit()),
		pool.WithWarmUpItems[*Session](cfg.PoolWarmUpSize()),
		pool.WithItemUsageLimit[*Session](cfg.PoolSessionUsageLimit()),
		pool.WithItemUsageTTL[*Session](cfg.PoolSessionUsageTTL()),
		pool.WithTrace[*Session](poolTrace(cfg.Trace())),
		pool.WithCreateItemTimeout[*Session](cfg.SessionCreateTimeout()),
		pool.WithCloseItemTimeout[*Session](cfg.SessionDeleteTimeout()),
		pool.WithMustDeleteItemFunc(func(s *Session, err error) bool {
			if !s.IsAlive() {
				return true
			}

			return err != nil && xerrors.MustDeleteTableOrQuerySession(err)
		}),
		pool.WithIdleTimeToLive[*Session](cfg.SessionIdleTimeToLive()),
		pool.WithCreateItemFunc(func(ctx context.Context) (_ *Session, err error) {
			if !cfg.DisableSessionBalancer() {
				ctx = meta.WithAllowFeatures(ctx, meta.HintSessionBalancer)
			}

			s, err := createSession(ctx, client,
				WithConn(cc),
				WithDeleteTimeout(cfg.SessionDeleteTimeout()),
				WithRegisterCloseCancel(c.registerCloseCancel),
				WithTrace(cfg.Trace()),
			)
			if err != nil {
				return nil, xerrors.WithStackTrace(err)
			}

			s.lazyTx = cfg.LazyTx()

			return s, nil
		}),
	)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	c.explicitSessionPool = explicitSessionPool

	implicitSessionPool, err := createImplicitSessionPool(ctx, cfg, client, cc)
	if err != nil {
		_ = c.explicitSessionPool.Close(ctx)

		return nil, xerrors.WithStackTrace(err)
	}

	c.implicitSessionPool = implicitSessionPool

	return c, nil
}

func poolTrace(t *trace.Query) *pool.Trace[*Session, Session] {
	return &pool.Trace[*Session, Session]{
		OnNew: func(ctx *context.Context, call stack.Caller) func(limit int) {
			onDone := gtrace.QueryOnPoolNew(t, ctx, call)

			return func(limit int) {
				onDone(limit)
			}
		},
		OnClose: func(ctx *context.Context, call stack.Caller) func(err error) {
			onDone := gtrace.QueryOnClose(t, ctx, call)

			return func(err error) {
				onDone(err)
			}
		},
		OnTry: func(ctx *context.Context, call stack.Caller) func(err error) {
			onDone := gtrace.QueryOnPoolTry(t, ctx, call)

			return func(err error) {
				onDone(err)
			}
		},
		OnWith: func(ctx *context.Context, call stack.Caller) func(attempts int, err error) {
			onDone := gtrace.QueryOnPoolWith(t, ctx, call)

			return func(attempts int, err error) {
				onDone(attempts, err)
			}
		},
		OnPut: func(ctx *context.Context, call stack.Caller, item *Session) func(err error) {
			onDone := gtrace.QueryOnPoolPut(t, ctx, call, safe.SessionInfo(item))

			return func(err error) {
				onDone(err)
			}
		},
		OnGet: func(ctx *context.Context, call stack.Caller) func(
			session *Session,
			hint *trace.NodeHintInfo,
			attempts int,
			err error,
		) {
			onDone := gtrace.QueryOnPoolGet(t, ctx, call)

			return func(session *Session, hint *trace.NodeHintInfo, attempts int, err error) {
				onDone(safe.SessionInfo(session), attempts, hint, err)
			}
		},
		OnChange: func(stats pool.Stats) {
			gtrace.QueryOnPoolChange(t, stats.Limit, stats.Idle, stats.CreateInProgress, stats.Concurrency, stats.Size)
		},
	}
}

func createImplicitSessionPool(ctx context.Context,
	cfg *config.Config,
	c Ydb_Query_V1.QueryServiceClient,
	cc grpc.ClientConnInterface,
) (sessionPool, error) {
	return pool.New(ctx,
		pool.WithLimit[*Session](cfg.PoolLimit()),
		pool.WithTrace[*Session](poolTrace(cfg.Trace())),
		pool.WithCreateItemFunc(func(ctx context.Context) (_ *Session, err error) {
			core := &implicitSessionCore{
				sessionCore{
					cc:     cc,
					Client: c,
					Trace:  cfg.Trace(),
				},
			}

			return &Session{
				Core:   core,
				trace:  cfg.Trace(),
				client: c,
			}, nil
		}),
	)
}
