package main

// postgres's admin client — postgres's own code, living in the module (novox/hq ADR 0039). Both this
// module's tools and its provisioner use it, and nothing outside postgres does.
//
// SQL goes over the wire protocol (pgconn), in the simple query protocol: a statement is sent as it
// was given, the way `psql -c` sent it when this module was TypeScript and could take no driver. One
// boundary, Dialer, and every method is built on it — which is also what the tests replace.

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/jackc/pgx/v5/pgconn"
)

// Reader is the login a caller's statement runs as (novox/hq issue 193). It may read every table and
// change nothing: `pg_read_all_data` and no other grant, and every transaction it opens is read-only
// by the server's own setting. A statement cannot climb out of a login the way it can out of a
// transaction wrapped around it as text: `COMMIT; DROP …` ended the old wrapper and ran the rest as
// the superuser, and even one read-only statement as a superuser can run a program on the server.
const Reader = "mesh_store_reader"

// readerOptions make the reader's session read-only from its first statement, before the role's own
// setting is read (PGOPTIONS, when this module shelled out to psql).
var readerOptions = map[string]string{
	"default_transaction_read_only": "on",
	"statement_timeout":             "60s",
}

// maxResultBytes bounds what one read-only query may hand back, as psql's 16 MiB buffer did.
const maxResultBytes = 16 << 20

// Result is one statement's answer: the command tag's verb and the rows, each a column's text or nil.
type Result struct {
	Command string
	Fields  []string
	Rows    [][]*string
}

// Maps is the rows keyed by their columns, as the tools return them.
func (r Result) Maps() []map[string]any {
	out := make([]map[string]any, 0, len(r.Rows))
	for _, row := range r.Rows {
		m := map[string]any{}
		for i, name := range r.Fields {
			if i < len(row) && row[i] != nil {
				m[name] = *row[i]
			} else {
				m[name] = nil
			}
		}
		out = append(out, m)
	}
	return out
}

// Session is one open connection: Run sends SQL in the simple protocol and answers every statement's
// result, in order.
type Session interface {
	Run(ctx context.Context, sql string) ([]Result, error)
	Close()
}

// Login is who a session connects as, to which database, with which session settings.
type Login struct {
	Database string
	User     string
	Password string
	Options  map[string]string
}

// Dialer opens a session. The real one is pgconn; a test's records what it was asked.
type Dialer func(ctx context.Context, l Login) (Session, error)

// Conn is where the server is and who administers it.
type Conn struct {
	Host     string
	Port     int
	User     string
	Password string
	SSLMode  string
	// ReaderPassword is the read-only login's password, which the mesh mints for this module
	// (`own-secrets.reader`). Empty when the mesh has not delivered it: then a caller's statement is
	// refused, never run as the admin (novox/hq issue 193).
	ReaderPassword string
}

// Client is postgres's admin client.
type Client struct {
	conn Conn
	dial Dialer

	readerMu    sync.Mutex
	readerReady bool
}

// NewClient is a client over a dialer; nil dials the server for real.
func NewClient(conn Conn, dial Dialer) *Client {
	if dial == nil {
		dial = pgDialer(conn)
	}
	return &Client{conn: conn, dial: dial}
}

var unfilled = regexp.MustCompile(`^\$\{[^}]*\}$`)

// ClientFromEnv builds the client from the module's words. MESH_POSTGRES_* first (the documented
// names), then the MESH_PROVISION_* keys the manifest sets. Fails without a host and an admin password.
func ClientFromEnv(env func(string) string) (*Client, error) {
	var u *url.URL
	if raw := env("MESH_PROVISION_POSTGRES"); raw != "" {
		if parsed, err := url.Parse(raw); err == nil {
			u = parsed
		}
	}
	host := env("MESH_POSTGRES_HOST")
	if host == "" && u != nil {
		host = u.Hostname()
	}
	// MESH_PROVISION_POSTGRES_PORT is the seat's twin (mesh-controller's internal/envfile.Placed
	// pattern): which port this machine actually put mesh-store at, when that differs from the
	// connection string's. Empty, or a placeholder the mesh never filled, adds nothing.
	seatPort := strings.TrimSpace(env("MESH_PROVISION_POSTGRES_PORT"))
	if unfilled.MatchString(seatPort) {
		seatPort = ""
	}
	portSource := firstOf(env("MESH_POSTGRES_PORT"), seatPort)
	if portSource == "" && u != nil {
		portSource = u.Port()
	}
	port, err := strconv.Atoi(portSource)
	if err != nil || port == 0 {
		port = 5432
	}
	user := env("MESH_POSTGRES_USER")
	if user == "" && u != nil && u.User != nil {
		user = u.User.Username()
	}
	if user == "" {
		user = "postgres"
	}
	password := firstOf(env("MESH_POSTGRES_PASSWORD"), readSecretFile(env("MESH_PROVISION_PASSWORD_FILE")))
	if host == "" || password == "" {
		return nil, errors.New("postgres host or admin password is not set — postgres's own code cannot reach the server")
	}
	sslmode := "prefer"
	if u != nil && u.Query().Get("sslmode") != "" {
		sslmode = u.Query().Get("sslmode")
	}
	reader := firstOf(env("MESH_POSTGRES_READER_PASSWORD"), readSecretFile(env("MESH_POSTGRES_READER_PASSWORD_FILE")))
	return NewClient(Conn{Host: host, Port: port, User: user, Password: password, SSLMode: sslmode, ReaderPassword: reader}, nil), nil
}

func firstOf(values ...string) string {
	for _, v := range values {
		if v != "" {
			return v
		}
	}
	return ""
}

func readSecretFile(path string) string {
	if path == "" {
		return ""
	}
	b, err := os.ReadFile(path)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// ---- the boundary -------------------------------------------------------------------------------

// as runs SQL as one login against one database and answers the last statement's result.
func (c *Client) as(ctx context.Context, l Login, sql string) (Result, error) {
	s, err := c.dial(ctx, l)
	if err != nil {
		return Result{}, err
	}
	defer s.Close()
	results, err := s.Run(ctx, sql)
	if err != nil {
		return Result{}, err
	}
	if len(results) == 0 {
		return Result{}, nil
	}
	// The last statement that answered rows, or the last statement: what psql -c printed last.
	for i := len(results) - 1; i >= 0; i-- {
		if len(results[i].Fields) > 0 {
			return results[i], nil
		}
	}
	return results[len(results)-1], nil
}

// Query runs SQL as the admin against a database ("postgres" when empty).
func (c *Client) Query(ctx context.Context, database, sql string) (Result, error) {
	if database == "" {
		database = "postgres"
	}
	return c.as(ctx, Login{Database: database, User: c.conn.User, Password: c.conn.Password}, sql)
}

func (c *Client) exists(ctx context.Context, sql string) (bool, error) {
	r, err := c.Query(ctx, "", sql)
	return len(r.Rows) > 0, err
}

// ---- what the provisioner does -----------------------------------------------------------------

// RoleStatement is the DDL that makes or re-sets a consumer's login. VALID UNTIL 'infinity': an
// expired password is refused like a wrong one, so the check the provisioner runs would report it
// lost, and only clearing the expiry makes applying it again work.
func RoleStatement(exists bool, role, password string) string {
	verb := "CREATE"
	if exists {
		verb = "ALTER"
	}
	return fmt.Sprintf("%s ROLE %s WITH LOGIN PASSWORD %s VALID UNTIL 'infinity'", verb, Ident(role), Literal(password))
}

// CreateDatabaseAndRole makes a login role and a database it owns, idempotently. A database that
// exists is never recreated, and nothing here drops anything.
func (c *Client) CreateDatabaseAndRole(ctx context.Context, database, role, password string) error {
	has, err := c.exists(ctx, "SELECT 1 FROM pg_roles WHERE rolname = "+Literal(role))
	if err != nil {
		return err
	}
	if _, err := c.Query(ctx, "", RoleStatement(has, role, password)); err != nil {
		return err
	}
	has, err = c.exists(ctx, "SELECT 1 FROM pg_database WHERE datname = "+Literal(database))
	if err != nil {
		return err
	}
	if !has {
		if _, err := c.Query(ctx, "", fmt.Sprintf("CREATE DATABASE %s OWNER %s", Ident(database), Ident(role))); err != nil {
			return err
		}
	}
	_, err = c.Query(ctx, "", fmt.Sprintf("GRANT ALL PRIVILEGES ON DATABASE %s TO %s", Ident(database), Ident(role)))
	return err
}

// Extensions reads a contribution's `extensions`: absent is none; otherwise a list of names.
func Extensions(values map[string]any) ([]string, error) {
	raw, ok := values["extensions"]
	if !ok || raw == nil {
		return nil, nil
	}
	list, ok := raw.([]any)
	if !ok {
		return nil, fmt.Errorf("extensions must be a list of extension names, not %T", raw)
	}
	seen := map[string]bool{}
	var out []string
	for _, v := range list {
		name, ok := v.(string)
		name = strings.TrimSpace(name)
		if !ok || name == "" {
			return nil, fmt.Errorf("extensions must be a list of extension names; %v is not one", v)
		}
		if !seen[name] {
			seen[name] = true
			out = append(out, name)
		}
	}
	return out, nil
}

// ExtensionStatement is the DDL that installs one extension in the database it is run in. IF NOT
// EXISTS: run on every pass, and a second run changes nothing. There is no statement that removes one.
func ExtensionStatement(name string) string {
	return "CREATE EXTENSION IF NOT EXISTS " + Ident(name)
}

// Unavailable is the names asked for that the server does not offer.
func Unavailable(want []string, available map[string]bool) []string {
	var missing []string
	for _, name := range want {
		if !available[name] {
			missing = append(missing, name)
		}
	}
	return missing
}

// EnsureExtensions installs each named extension in the database, as the admin, connected to that
// database — most extensions (pgvector among them) are not trusted, so the consumer that owns the
// database cannot install them itself. Only names the server lists in pg_available_extensions are
// asked for; any other is refused, by name, before anything runs. An extension is never dropped.
func (c *Client) EnsureExtensions(ctx context.Context, database string, want []string) error {
	if len(want) == 0 {
		return nil
	}
	r, err := c.Query(ctx, "", "SELECT name FROM pg_available_extensions")
	if err != nil {
		return err
	}
	available := map[string]bool{}
	for _, row := range r.Rows {
		if len(row) > 0 && row[0] != nil {
			available[*row[0]] = true
		}
	}
	if missing := Unavailable(want, available); len(missing) > 0 {
		return fmt.Errorf("database %s asks for extension(s) %s, which this server does not offer "+
			"(not in pg_available_extensions); refused, nothing installed", database, strings.Join(quoted(missing), ", "))
	}
	for _, name := range want {
		if _, err := c.Query(ctx, database, ExtensionStatement(name)); err != nil {
			return fmt.Errorf("extension %q in %s: %w", name, database, err)
		}
	}
	return nil
}

func quoted(names []string) []string {
	out := make([]string, len(names))
	for i, n := range names {
		out[i] = strconv.Quote(n)
	}
	return out
}

// lostCodes are the server's ways of saying a login, its password or its database is wrong or gone:
// invalid_password, invalid_authorization_specification (no such role, not permitted to log in),
// invalid_catalog_name (no such database), insufficient_privilege (no CONNECT).
var lostCodes = map[string]bool{"28P01": true, "28000": true, "3D000": true, "42501": true}

// IsLost says whether an error is the server saying the credential is wrong or gone.
func IsLost(err error) bool {
	var pg *pgconn.PgError
	return errors.As(err, &pg) && lostCodes[pg.Code]
}

// Holds says whether `role` can log in to `database` with exactly `password` and finds every wanted
// extension installed there: the consumer's own view, checked by connecting as it. Read-only. false
// only when the server says so; an unreachable server is an error, because being unable to ask is
// not evidence of loss (novox/hq issue 120).
func (c *Client) Holds(ctx context.Context, database, role, password string, extensions []string) (bool, error) {
	ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()
	r, err := c.as(ctx, Login{Database: database, User: role, Password: password}, "SELECT extname FROM pg_extension")
	if err != nil {
		if IsLost(err) {
			return false, nil
		}
		return false, err
	}
	installed := map[string]bool{}
	for _, row := range r.Rows {
		if len(row) > 0 && row[0] != nil {
			installed[*row[0]] = true
		}
	}
	return len(Unavailable(extensions, installed)) == 0, nil
}

// LockRole withdraws a consumer without destroying anything (novox/hq issue 241): its login can no
// longer log in and its open connections are ended, and its database stays exactly as it was, under
// its own name. A consumer that comes back is given the same database — create sets LOGIN again.
func (c *Client) LockRole(ctx context.Context, role string) error {
	has, err := c.exists(ctx, "SELECT 1 FROM pg_roles WHERE rolname = "+Literal(role))
	if err != nil || !has {
		return err
	}
	if _, err := c.Query(ctx, "", fmt.Sprintf("ALTER ROLE %s NOLOGIN", Ident(role))); err != nil {
		return err
	}
	_, err = c.Query(ctx, "", "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE usename = "+
		Literal(role)+" AND pid <> pg_backend_pid()")
	return err
}

// ---- what the tools do --------------------------------------------------------------------------

// RetireDatabase takes a database out of service on purpose: renamed aside to
// `<name>_deleted_<date>` and its owner locked. Never a drop — the data stays on the server under the
// new name until a person removes it by hand. Answers the name it now has.
func (c *Client) RetireDatabase(ctx context.Context, database string, now time.Time) (string, error) {
	found, err := c.Query(ctx, "", "SELECT pg_get_userbyid(datdba) AS owner FROM pg_database WHERE datname = "+Literal(database))
	if err != nil {
		return "", err
	}
	if len(found.Rows) == 0 {
		return "", fmt.Errorf("no database named %s", database)
	}
	owner := ""
	if row := found.Rows[0]; len(row) > 0 && row[0] != nil {
		owner = *row[0]
	}
	aside := RetiredName(database, now)
	taken, err := c.exists(ctx, "SELECT 1 FROM pg_database WHERE datname = "+Literal(aside))
	if err != nil {
		return "", err
	}
	if taken {
		return "", fmt.Errorf("%s already exists; retire it by hand first", aside)
	}
	if owner != "" && owner != "postgres" {
		if _, err := c.Query(ctx, "", fmt.Sprintf("ALTER ROLE %s NOLOGIN", Ident(owner))); err != nil {
			return "", err
		}
	}
	if _, err := c.Query(ctx, "", "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = "+
		Literal(database)+" AND pid <> pg_backend_pid()"); err != nil {
		return "", err
	}
	if _, err := c.Query(ctx, "", fmt.Sprintf("ALTER DATABASE %s RENAME TO %s", Ident(database), Ident(aside))); err != nil {
		return "", err
	}
	return aside, nil
}

// Database is one row of the listing.
type Database struct {
	Name      string `json:"name"`
	SizeBytes int64  `json:"sizeBytes"`
}

// ListDatabases is the non-template databases with their size.
func (c *Client) ListDatabases(ctx context.Context) ([]Database, error) {
	r, err := c.Query(ctx, "", "SELECT datname, pg_database_size(datname) AS size FROM pg_database WHERE datistemplate = false ORDER BY datname")
	if err != nil {
		return nil, err
	}
	out := []Database{}
	for _, row := range r.Rows {
		if len(row) < 2 || row[0] == nil {
			continue
		}
		d := Database{Name: *row[0]}
		if row[1] != nil {
			d.SizeBytes, _ = strconv.ParseInt(*row[1], 10, 64)
		}
		out = append(out, d)
	}
	return out, nil
}

// ReaderStatements are the DDL that make the read-only login, every attribute stated so an existing
// role someone widened is narrowed again.
func ReaderStatements(exists bool, password string) []string {
	verb := "CREATE"
	if exists {
		verb = "ALTER"
	}
	return []string{
		fmt.Sprintf("%s ROLE %s WITH LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION "+
			"NOBYPASSRLS INHERIT PASSWORD %s VALID UNTIL 'infinity'", verb, Ident(Reader), Literal(password)),
		"GRANT pg_read_all_data TO " + Ident(Reader),
		"ALTER ROLE " + Ident(Reader) + " SET default_transaction_read_only = on",
		"ALTER ROLE " + Ident(Reader) + " SET statement_timeout = '60s'",
	}
}

var errReaderMissing = errors.New("the read-only login's password was not delivered (own-secrets.reader, " +
	"MESH_POSTGRES_READER_PASSWORD_FILE), so the statement is refused rather than run as the admin (novox/hq issue 193)")

// EnsureReader makes the read-only login, idempotently, with the password the mesh minted — as the
// admin, because only the admin can make a role. Once per process; a failure is asked again next call.
func (c *Client) EnsureReader(ctx context.Context) error {
	if c.conn.ReaderPassword == "" {
		return errReaderMissing
	}
	c.readerMu.Lock()
	defer c.readerMu.Unlock()
	if c.readerReady {
		return nil
	}
	has, err := c.exists(ctx, "SELECT 1 FROM pg_roles WHERE rolname = "+Literal(Reader))
	if err != nil {
		return err
	}
	for _, sql := range ReaderStatements(has, c.conn.ReaderPassword) {
		if _, err := c.Query(ctx, "", sql); err != nil {
			return err
		}
	}
	c.readerReady = true
	return nil
}

var firstWord = regexp.MustCompile(`^\s*([A-Za-z]+)`)

// ReadOnlyQuery runs a caller's statement against a named database as the read-only login, for the
// postgres_query tool and the store seat's `query` verb (novox/hq ADR 0159, issue 193).
//
// **Read-only by the login, not by text around the statement.** The statement is sent as it was
// given, as the reader, whose role can write nothing and whose session the server makes read-only.
// Never as the admin: without the reader's password the call is refused.
func (c *Client) ReadOnlyQuery(ctx context.Context, database, sql string) (Result, error) {
	if err := c.EnsureReader(ctx); err != nil {
		return Result{}, err
	}
	r, err := c.as(ctx, Login{Database: database, User: Reader, Password: c.conn.ReaderPassword, Options: readerOptions}, sql)
	if err != nil {
		return Result{}, err
	}
	size := 0
	for _, row := range r.Rows {
		for _, v := range row {
			if v != nil {
				size += len(*v)
			}
		}
	}
	if size > maxResultBytes {
		return Result{}, fmt.Errorf("the result is larger than %d MiB; narrow the query", maxResultBytes>>20)
	}
	r.Command = ""
	if m := firstWord.FindStringSubmatch(sql); m != nil {
		r.Command = strings.ToUpper(m[1])
	}
	return r, nil
}

// ---- quoting and names -------------------------------------------------------------------------

// Ident quotes a SQL identifier: double quotes, internal ones doubled.
func Ident(id string) string { return `"` + strings.ReplaceAll(id, `"`, `""`) + `"` }

// Literal quotes a SQL string literal: single quotes, internal ones doubled (standard_conforming_strings).
func Literal(v string) string { return "'" + strings.ReplaceAll(v, "'", "''") + "'" }

// RetiredName is `<name>_deleted_<yyyymmdd>`, within postgres's 63 bytes.
func RetiredName(database string, now time.Time) string {
	suffix := "_deleted_" + now.UTC().Format("20060102")
	keep := 63 - len(suffix)
	if len(database) > keep {
		database = database[:keep]
		for !utf8.ValidString(database) {
			database = database[:len(database)-1]
		}
	}
	return database + suffix
}

// ---- the real dialer ---------------------------------------------------------------------------

type pgSession struct{ c *pgconn.PgConn }

func (s pgSession) Run(ctx context.Context, sql string) ([]Result, error) {
	results, err := s.c.Exec(ctx, sql).ReadAll()
	if err != nil {
		return nil, err
	}
	out := make([]Result, 0, len(results))
	for _, r := range results {
		if r.Err != nil {
			return nil, r.Err
		}
		res := Result{Command: strings.SplitN(r.CommandTag.String(), " ", 2)[0]}
		for _, f := range r.FieldDescriptions {
			res.Fields = append(res.Fields, f.Name)
		}
		for _, row := range r.Rows {
			cells := make([]*string, len(row))
			for i, v := range row {
				if v != nil {
					s := string(v)
					cells[i] = &s
				}
			}
			res.Rows = append(res.Rows, cells)
		}
		out = append(out, res)
	}
	return out, nil
}

func (s pgSession) Close() {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = s.c.Close(ctx)
}

func pgDialer(conn Conn) Dialer {
	return func(ctx context.Context, l Login) (Session, error) {
		u := url.URL{
			Scheme:   "postgres",
			User:     url.UserPassword(l.User, l.Password),
			Host:     net.JoinHostPort(conn.Host, strconv.Itoa(conn.Port)),
			Path:     "/" + l.Database,
			RawQuery: url.Values{"sslmode": {firstOf(conn.SSLMode, "prefer")}, "connect_timeout": {"10"}}.Encode(),
		}
		cfg, err := pgconn.ParseConfig(u.String())
		if err != nil {
			return nil, err
		}
		for k, v := range l.Options {
			cfg.RuntimeParams[k] = v
		}
		c, err := pgconn.ConnectConfig(ctx, cfg)
		if err != nil {
			return nil, err
		}
		return pgSession{c: c}, nil
	}
}