package postgres

import (
	"context"
	"database/sql"
	"net"
	"net/url"
	"strings"

	"code.crute.us/mcrute/golib/secrets"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/jackc/pgx/v5/stdlib"
)

func prepareUri(ctx context.Context, uri string, sc secrets.Client) (string, string, error) {
	// Prefix uri with postgres:// unless it already includes one of the
	// standard prefixes. Otherwise if scheme is omitted then url parsing
	// will fail to capture the username for secret lookup.
	if !strings.HasPrefix(uri, "postgres://") {
		uri = "postgres://" + uri
	}

	u, err := url.Parse(uri)
	if err != nil {
		return "", "", err
	}

	// The username provided by the user (there should be no
	// password) will be a reference to a secret material with the
	// prefix database/creds/. This needs to be replaced with the real
	// username/password pair fetched from secret before attempting to
	// connect.
	cred, _, err := sc.DatabaseCredential(ctx, u.User.Username())
	if err != nil {
		return "", "", err
	}
	u.User = url.UserPassword(cred.Username, cred.Password)

	// User may imply the default port
	if u.Port() == "" {
		u.Host = net.JoinHostPort(u.Host, "5432")
	}

	return u.String(), u.Query().Get("TimeZone"), nil
}

func Connect(ctx context.Context, uri string, sc secrets.Client) (*pgx.Conn, error) {
	dsn, tz, err := prepareUri(ctx, uri, sc)
	if err != nil {
		return nil, err
	}

	cfg, err := pgx.ParseConfig(dsn)
	if err != nil {
		return nil, err
	}

	if tz != "" {
		cfg.RuntimeParams["timezone"] = tz
	}

	return pgx.ConnectConfig(ctx, cfg)
}

func ConnectPool(ctx context.Context, uri string, sc secrets.Client) (*pgxpool.Pool, error) {
	dsn, tz, err := prepareUri(ctx, uri, sc)
	if err != nil {
		return nil, err
	}

	cfg, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, err
	}

	if tz != "" {
		cfg.ConnConfig.RuntimeParams["timezone"] = tz
	}

	return pgxpool.NewWithConfig(ctx, cfg)
}

func ConnectStdlibPool(ctx context.Context, uri string, sc secrets.Client) (*sql.DB, error) {
	pool, err := ConnectPool(ctx, uri, sc)
	if err != nil {
		return nil, err
	}

	return stdlib.OpenDBFromPool(pool), nil
}