package postgres

import (
	"context"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pkg/errors"
)

// From https://gitlab.ozon.dev/go/classroom-18/students/week-4-workshop/-/blob/master/internal/infra/postgres/postgres.go

func NewPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
	config, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, errors.Wrap(err, "pgxpool.ParseConfig")
	}

	config.ConnConfig.Tracer = promTracer{}

	pool, err := pgxpool.NewWithConfig(ctx, config)
	if err != nil {
		return nil, errors.Wrap(err, "pgxpool.NewWithConfig")
	}

	return pool, nil
}

func NewPools(ctx context.Context, DSNs ...string) ([]*pgxpool.Pool, error) {
	pools := make([]*pgxpool.Pool, len(DSNs))

	for i, dsn := range DSNs {
		cfg, err := pgxpool.ParseConfig(dsn)
		if err != nil {
			closeOpened(pools[:i])

			return nil, errors.Wrap(err, "pgxpool.ParseConfig")
		}

		cfg.ConnConfig.Tracer = promTracer{}

		pool, err := pgxpool.NewWithConfig(ctx, cfg)
		if err != nil {
			closeOpened(pools[:i])

			return nil, errors.Wrap(err, "pgxpool.NewWithConfig")
		}

		pools[i] = pool
	}

	return pools, nil
}

func closeOpened(pools []*pgxpool.Pool) {
	for _, p := range pools {
		if p != nil {
			p.Close()
		}
	}
}