// Copyright 2019 The Go Cloud Development Kit Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package postgres provides functions to open PostgreSQL databases with OpenTelemetry instrumentation.
package postgres

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

	"github.com/XSAM/otelsql"
	"github.com/lib/pq"
	"gocloud.dev/internal/openurl"
)

// Scheme is the URL scheme this package registers its URLOpener under on
// DefaultMux.
const Scheme = "postgres"

func init() {
	DefaultURLMux().RegisterPostgres(Scheme, &URLOpener{})
}

// URLOpener opens URLs like "postgres://" by using the underlying PostgreSQL driver.
// See https://godoc.org/github.com/lib/pq#hdr-Connection_String_Parameters for details.
type URLOpener struct {
	TraceOpts []otelsql.Option
}

// OpenPostgresURL opens a new database connection wrapped with OpenTelemetry instrumentation.
func (uo *URLOpener) OpenPostgresURL(ctx context.Context, u *url.URL) (*sql.DB, error) {
	u2 := new(url.URL)
	*u2 = *u
	u2.Scheme = "postgres"
	return sql.OpenDB(connector{dsn: u2.String(), traceOpts: append([]otelsql.Option(nil), uo.TraceOpts...)}), nil
}

type connector struct {
	dsn       string
	traceOpts []otelsql.Option
}

func (c connector) Connect(ctx context.Context) (driver.Conn, error) {
	return c.Driver().Open(c.dsn)
}

func (c connector) Driver() driver.Driver {
	return otelsql.WrapDriver(&pq.Driver{}, c.traceOpts...)
}

// PostgresURLOpener can open Postgres connections based on a URL.
// The opener must not modify the URL argument. OpenPostgresURL must be safe to
// call from multiple goroutines.
//
// This interface is generally implemented by types in driver packages.
type PostgresURLOpener interface {
	// OpenPostgresURL opens a Postgres connection based on the URL.
	OpenPostgresURL(ctx context.Context, u *url.URL) (*sql.DB, error)
}

// URLMux is a URL opener multiplexer. It matches the scheme of the URLs
// against a set of registered schemes and calls the opener that matches the
// URL's scheme.
//
// The zero value is a multiplexer with no registered schemes.
type URLMux struct {
	schemes openurl.SchemeMap
}

// RegisterPostgres registers the opener with the given scheme. If an opener
// already exists for the scheme, RegisterPostgres panics.
func (mux *URLMux) RegisterPostgres(scheme string, opener PostgresURLOpener) {
	mux.schemes.Register("postgres", "DB", scheme, opener)
}

// OpenPostgres calls OpenPostgresURL with the URL parsed from urlstr.
// OpenPostgres is safe to call from multiple goroutines.
func (mux *URLMux) OpenPostgres(ctx context.Context, urlstr string) (*sql.DB, error) {
	opener, u, err := mux.schemes.FromString("DB", urlstr)
	if err != nil {
		return nil, err
	}
	return opener.(PostgresURLOpener).OpenPostgresURL(ctx, u)
}

// OpenPostgresURL dispatches the URL to the opener that is registered with the
// URL's scheme. OpenPostgresURL is safe to call from multiple goroutines.
func (mux *URLMux) OpenPostgresURL(ctx context.Context, u *url.URL) (*sql.DB, error) {
	opener, err := mux.schemes.FromURL("DB", u)
	if err != nil {
		return nil, err
	}
	return opener.(PostgresURLOpener).OpenPostgresURL(ctx, u)
}

var defaultURLMux = new(URLMux)

// DefaultURLMux returns the URLMux used by OpenPostgres.
//
// Driver packages can use this to register their PostgresURLOpener on the mux.
func DefaultURLMux() *URLMux {
	return defaultURLMux
}

// Open opens the bucket identified by the URL given. URL openers must be
// registered in the DefaultURLMux, which is typically done in driver
// packages' initialization.
//
// See the URLOpener documentation in driver subpackages for more
// details on supported scheme(s) and URL parameter(s).
func Open(ctx context.Context, urlstr string) (*sql.DB, error) {
	return defaultURLMux.OpenPostgres(ctx, urlstr)
}