package clusterdb

import (
	"crypto/md5"
	"database/sql"
	"fmt"

	"code.justin.tv/d8a/buddy/lib/store"
)

//PostgresQueriesFactory is a DriverQueriesFactory implementation that just returns a PostgresQueries object
type PostgresQueriesFactory struct{}

//BuildDriverQueries just constructs and returns a PostgresQueries object
func (factory *PostgresQueriesFactory) BuildDriverQueries(db *sql.DB) (DriverQueries, error) {
	return &PostgresQueries{}, nil
}

// PostgresQueries is a DriverQueries implementation for any postgres database
type PostgresQueries struct{}

//CanLockUsers returns whether this backing database supports user lock/unlock to prevent logins for a particular user
func (queries *PostgresQueries) CanLockUsers() bool {
	return true
}

//CanUseRolePairs returns whether this backing database supports HA password rotations with role pairs
func (queries *PostgresQueries) CanUseRolePairs() bool {
	return true
}

// TestConnection tests the DB connection by sending a SELECT 1, because sometimes the connection process
// itself doesn't error.
func (queries *PostgresQueries) TestConnection(db *sql.DB) error {
	rows, err := db.Query("SELECT 1")
	if rows != nil {
		store.TryClose(rows)
	}

	return err
}

// AlterPassword modifies the DB user's password in the clusterDB
func (queries *PostgresQueries) AlterPassword(db *sql.DB, username string, password string) (sql.Result, error) {
	encryptedPass := md5.Sum([]byte(password + username))
	return db.Exec(fmt.Sprintf("ALTER USER \"%s\" WITH ENCRYPTED PASSWORD 'md5%x'", username, encryptedPass))
}

// SetUserLock sets the DB user to locked or unlocked for logins
func (queries *PostgresQueries) SetUserLock(db *sql.DB, username string, locked bool) (sql.Result, error) {
	newValue := "LOGIN"
	if locked {
		newValue = "NOLOGIN"
	}

	return db.Exec(fmt.Sprintf("ALTER ROLE \"%s\" WITH %s", username, newValue))
}

// CreateRole makes a new blank DB role in the cluster db
func (queries *PostgresQueries) CreateRole(db *sql.DB, roleName string) (sql.Result, error) {
	return db.Exec(fmt.Sprintf("CREATE ROLE \"%s\" NOLOGIN", roleName))
}

// CreateUser makes a new DB user in the cluster db, with password and an inherited role
func (queries *PostgresQueries) CreateUser(db *sql.DB, username string, password string, parentRole string) (sql.Result, error) {
	encryptedPass := md5.Sum([]byte(password + username))
	return db.Exec(fmt.Sprintf("CREATE ROLE \"%s\" LOGIN ENCRYPTED PASSWORD 'md5%x' IN ROLE \"%s\"", username, encryptedPass, parentRole))
}

// RoleParents returns a list of role inheritance relationships in the clusterDb, as a map
// of role name -> list of member role names
func (queries *PostgresQueries) RoleParents(db *sql.DB) (*sql.Rows, error) {
	return db.Query("select role.rolname AS role, usr.rolname AS user FROM pg_auth_members members INNER JOIN pg_roles AS role ON members.roleid=role.oid INNER JOIN pg_roles AS usr ON members.member=usr.oid")
}

// GrantedRoles returns a list of role names who have grants in the cluster DB
func (queries *PostgresQueries) GrantedRoles(db *sql.DB) (*sql.Rows, error) {
	return db.Query(`
		select DISTINCT(grantee) FROM information_schema.role_column_grants 
			UNION 
		select distinct(grantee) FROM information_schema.role_udt_grants 
			UNION 
		select distinct(grantee) FROM information_schema.role_table_grants 
			UNION 
		select distinct(grantee) FROM information_schema.role_routine_grants 
			UNION 
		select distinct(grantee) FROM information_schema.role_usage_grants`)
}

// ClusterUsers retrieves all users in the cluster DB, as a map of username -> locked/unlocked  If true, unlocked, if false locked
func (queries *PostgresQueries) ClusterUsers(db *sql.DB) (*sql.Rows, error) {
	return db.Query("SELECT rolname, rolcanlogin FROM pg_catalog.pg_roles WHERE NOT rolsuper")
}

// ActiveUsers returns a list of users currently executing a query on the server
func (queries *PostgresQueries) ActiveUsers(db *sql.DB) (*sql.Rows, error) {
	return db.Query("SELECT DISTINCT(usename) FROM pg_stat_activity")
}
