package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"os"

	_ "github.com/go-sql-driver/mysql"
)

var wipeDB bool

func init() {
	flag.BoolVar(&wipeDB, "wipe-databases", false, "wipe databases before creating them")
}

func main() {
	flag.Parse()

	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:3306)/",
		os.Getenv("DB_ADMIN_USER"),
		os.Getenv("DB_ADMIN_PASS"),
		os.Getenv("DB_HOST"),
	))
	if err != nil {
		log.Fatal("failed to connect to database:", err)
	}

	if err := createDB(db,
		os.Getenv("DB_PASS"),
	); err != nil {
		log.Fatal("failed to ensure database creation:", err)
	}

	log.Print("database verified")
}

func createDB(db *sql.DB, dbpass string) error {
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("failed to create db transaction: %w", err)
	}

	defer tx.Rollback()

	if wipeDB {
		if _, err := tx.Exec("DROP DATABASE IF EXISTS mcp"); err != nil {
			return fmt.Errorf("failed to drop existing mcp database: %w", err)
		}

		if _, err := tx.Exec("DROP DATABASE IF EXISTS ss"); err != nil {
			return fmt.Errorf("failed to drop existing ss database: %w", err)
		}
	}

	// NB: "CREATE USER" does not support substitutions:  https://bugs.mysql.com/bug.php?id=28406
	if _, err := tx.Exec(fmt.Sprintf(`CREATE USER IF NOT EXISTS 'am'@'%%' IDENTIFIED BY '%s'`, dbpass)); err != nil {
		return fmt.Errorf("failed to create database user: %w", err)
	}

	if _, err := tx.Exec(`CREATE DATABASE IF NOT EXISTS mcp CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci`); err != nil {
		return fmt.Errorf("failed to ensure creation of mcp database: %w", err)
	}

	if _, err := tx.Exec(`GRANT ALL ON mcp.* TO 'am'@'%'`); err != nil {
		return fmt.Errorf("failed to ensure privilege to mcp database: %w", err)
	}

	if _, err := tx.Exec(`CREATE DATABASE IF NOT EXISTS ss CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci`); err != nil {
		return fmt.Errorf("failed to ensure creation of ss database: %w", err)
	}

	if _, err := tx.Exec(`GRANT ALL ON ss.* TO 'am'@'%'`); err != nil {
		return fmt.Errorf("failed to ensure privilege to ss database: %w", err)
	}

	if _, err := tx.Exec(`FLUSH PRIVILEGES`); err != nil {
		return fmt.Errorf("failed to flush privileges: %w", err)
	}

	return tx.Commit()
}