package main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
)

// cmdInitDB creates the PostgreSQL role and database for edgeguard if
// they don't already exist. Idempotent — safe to re-run from postinst
// on every package upgrade.
//
// Authentication: the local postgres super-user via unix-socket peer
// auth. edgeguard-ctl initdb is invoked from postinst (running as
// root), so it shells out via `sudo -u postgres psql -tA -c ...`.
//
// The created role has LOGIN but no password — production access is
// peer-auth from the system user `edgeguard` over the unix socket
// (architecture.md §6).
//
// Flags:
//
//	--db <name>    database name (default: edgeguard)
//	--role <name>  role name     (default: edgeguard)
func cmdInitDB(args []string) int {
	dbName := "edgeguard"
	roleName := "edgeguard"
	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "--db":
			if i+1 < len(args) {
				dbName = args[i+1]
				i++
			}
		case "--role":
			if i+1 < len(args) {
				roleName = args[i+1]
				i++
			}
		case "-h", "--help":
			fmt.Println("Usage: edgeguard-ctl initdb [--db NAME] [--role NAME]")
			return 0
		}
	}

	if !looksLikeIdentifier(dbName) || !looksLikeIdentifier(roleName) {
		fmt.Fprintln(os.Stderr, "initdb: --db and --role must match [a-z_][a-z0-9_]{0,62}")
		return 2
	}

	roleExists, err := psqlBoolQuery(fmt.Sprintf(
		"SELECT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '%s')", roleName))
	if err != nil {
		fmt.Fprintln(os.Stderr, "initdb: check role:", err)
		return 1
	}
	if roleExists {
		fmt.Printf("role %q already exists — leaving it\n", roleName)
	} else {
		if err := psqlExec(fmt.Sprintf("CREATE ROLE %s LOGIN", roleName)); err != nil {
			fmt.Fprintln(os.Stderr, "initdb: create role:", err)
			return 1
		}
		fmt.Printf("created role %q\n", roleName)
	}

	dbExists, err := psqlBoolQuery(fmt.Sprintf(
		"SELECT EXISTS (SELECT 1 FROM pg_database WHERE datname = '%s')", dbName))
	if err != nil {
		fmt.Fprintln(os.Stderr, "initdb: check database:", err)
		return 1
	}
	if dbExists {
		fmt.Printf("database %q already exists — leaving it\n", dbName)
	} else {
		if err := psqlExec(fmt.Sprintf("CREATE DATABASE %s OWNER %s", dbName, roleName)); err != nil {
			fmt.Fprintln(os.Stderr, "initdb: create database:", err)
			return 1
		}
		fmt.Printf("created database %q owned by %q\n", dbName, roleName)
	}

	return 0
}

// looksLikeIdentifier guards against shell- and SQL-injection in the
// role/db flags. Only allow the standard PG identifier shape; we
// interpolate into psql -c so a quoted user-supplied string is
// hostile.
func looksLikeIdentifier(s string) bool {
	if s == "" || len(s) > 63 {
		return false
	}
	if !(s[0] == '_' || (s[0] >= 'a' && s[0] <= 'z')) {
		return false
	}
	for _, r := range s[1:] {
		ok := r == '_' ||
			(r >= 'a' && r <= 'z') ||
			(r >= '0' && r <= '9')
		if !ok {
			return false
		}
	}
	return true
}

// psqlBoolQuery runs a one-shot SELECT EXISTS query as the postgres
// super-user and returns true/false. Tuples-only + unaligned mode
// (-tA) so we get just the literal "t" or "f".
func psqlBoolQuery(sql string) (bool, error) {
	out, err := psqlRun([]string{"-tA", "-c", sql})
	if err != nil {
		return false, err
	}
	return strings.TrimSpace(string(out)) == "t", nil
}

// psqlExec runs a CREATE ROLE / CREATE DATABASE statement as the
// postgres super-user. ON_ERROR_STOP makes psql return non-zero on
// any SQL error.
func psqlExec(sql string) error {
	if _, err := psqlRun([]string{"-v", "ON_ERROR_STOP=1", "-c", sql}); err != nil {
		return err
	}
	return nil
}

// psqlRun shells out to `sudo -u postgres psql` with the given args.
// When edgeguard-ctl initdb is already running as the postgres user
// (e.g. an admin re-running it manually), skip sudo so it works
// without sudoers configuration.
func psqlRun(args []string) ([]byte, error) {
	cmd := buildPsqlCmd(args)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	out, err := cmd.Output()
	if err != nil {
		msg := strings.TrimSpace(stderr.String())
		if msg != "" {
			return nil, fmt.Errorf("%w: %s", err, msg)
		}
		return nil, err
	}
	return out, nil
}

func buildPsqlCmd(args []string) *exec.Cmd {
	if u := os.Getenv("USER"); u == "postgres" {
		return exec.Command("psql", args...)
	}
	full := append([]string{"-n", "-u", "postgres", "psql"}, args...)
	return exec.Command("sudo", full...)
}