package database

import (
	"context"
	"crypto/sha256"
	"embed"
	"errors"
	"fmt"
	"io/fs"
	"sort"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"
)

const migrationLockID int64 = 648093714083

//go:embed migrations/*.up.sql
var migrations embed.FS

type MigrationOptions struct {
	DatabaseURL  string
	RuntimeRole  string
	IdentityRole string
	LockTimeout  time.Duration
}

type migrationFile struct {
	Name     string
	Checksum string
	Content  []byte
}

func RunMigrations(
	ctx context.Context,
	options MigrationOptions,
	logger *zap.Logger,
) error {
	if err := options.validate(); err != nil {
		return err
	}

	poolConfig, err := pgxpool.ParseConfig(options.DatabaseURL)
	if err != nil {
		return fmt.Errorf("parse migration database URL: %w", err)
	}
	poolConfig.ConnConfig.RuntimeParams["application_name"] = "ariadne-migrate"

	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
	if err != nil {
		return fmt.Errorf("open migration database: %w", err)
	}
	defer pool.Close()

	connection, err := pool.Acquire(ctx)
	if err != nil {
		return fmt.Errorf("acquire migration connection: %w", err)
	}
	defer connection.Release()

	if err := connection.Conn().Ping(ctx); err != nil {
		return fmt.Errorf("ping migration database: %w", err)
	}
	if err := validateMigrationRole(ctx, connection, options); err != nil {
		return err
	}
	if err := acquireMigrationLock(ctx, connection, options.LockTimeout); err != nil {
		return err
	}
	defer releaseMigrationLock(connection)

	if err := applyMigrations(ctx, connection); err != nil {
		return err
	}
	if err := synchronizeAccess(ctx, connection, options); err != nil {
		return err
	}

	logger.Info(
		"database migrations and access policies applied",
		zap.String("runtime_role", options.RuntimeRole),
		zap.String("identity_role", options.IdentityRole),
	)
	return nil
}

func (options MigrationOptions) validate() error {
	if strings.TrimSpace(options.DatabaseURL) == "" {
		return errors.New("ARIADNE_DATABASE_MIGRATION_URL is required")
	}
	if strings.TrimSpace(options.RuntimeRole) == "" {
		return errors.New("ARIADNE_DATABASE_RUNTIME_ROLE is required")
	}
	if strings.TrimSpace(options.IdentityRole) == "" {
		return errors.New("ARIADNE_DATABASE_IDENTITY_ROLE is required")
	}
	if options.RuntimeRole == options.IdentityRole {
		return errors.New("runtime and identity database roles must be different")
	}
	if options.LockTimeout <= 0 {
		return errors.New("ARIADNE_MIGRATION_LOCK_TIMEOUT must be greater than zero")
	}
	return nil
}

func validateMigrationRole(
	ctx context.Context,
	connection *pgxpool.Conn,
	options MigrationOptions,
) error {
	var (
		roleName      string
		isSuperuser   bool
		canBypassRLS  bool
		canCreateDB   bool
		canCreateRole bool
		canReplicate  bool
		canInherit    bool
	)
	if err := connection.QueryRow(ctx, `
		SELECT
			rolname,
			rolsuper,
			rolbypassrls,
			rolcreatedb,
			rolcreaterole,
			rolreplication,
			rolinherit
		FROM pg_roles
		WHERE rolname = current_user
	`).Scan(
		&roleName,
		&isSuperuser,
		&canBypassRLS,
		&canCreateDB,
		&canCreateRole,
		&canReplicate,
		&canInherit,
	); err != nil {
		return fmt.Errorf("inspect migration role: %w", err)
	}
	if isSuperuser ||
		!canBypassRLS ||
		canCreateDB ||
		canCreateRole ||
		canReplicate ||
		canInherit {
		return fmt.Errorf(
			"migration role %q must be NOSUPERUSER, BYPASSRLS, "+
				"NOCREATEDB, NOCREATEROLE, NOREPLICATION and NOINHERIT",
			roleName,
		)
	}
	if roleName == options.RuntimeRole || roleName == options.IdentityRole {
		return fmt.Errorf("migration role %q must be dedicated", roleName)
	}
	return nil
}

func acquireMigrationLock(
	ctx context.Context,
	connection *pgxpool.Conn,
	timeout time.Duration,
) error {
	lockContext, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	ticker := time.NewTicker(250 * time.Millisecond)
	defer ticker.Stop()

	for {
		var acquired bool
		if err := connection.QueryRow(
			lockContext,
			"SELECT pg_try_advisory_lock($1)",
			migrationLockID,
		).Scan(&acquired); err != nil {
			return fmt.Errorf("try migration lock: %w", err)
		}
		if acquired {
			return nil
		}

		select {
		case <-lockContext.Done():
			return fmt.Errorf(
				"acquire migration lock within %s: %w",
				timeout,
				lockContext.Err(),
			)
		case <-ticker.C:
		}
	}
}

func releaseMigrationLock(connection *pgxpool.Conn) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _ = connection.Exec(ctx, "SELECT pg_advisory_unlock($1)", migrationLockID)
}

func applyMigrations(ctx context.Context, connection *pgxpool.Conn) error {
	if _, err := connection.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS public.ariadne_schema_migrations (
			version text PRIMARY KEY,
			checksum text NOT NULL,
			applied_at timestamptz NOT NULL DEFAULT now()
		)
	`); err != nil {
		return fmt.Errorf("create migration table: %w", err)
	}

	manifest, err := embeddedMigrationManifest()
	if err != nil {
		return err
	}

	for _, migration := range manifest {
		var appliedChecksum string
		err := connection.QueryRow(
			ctx,
			`
				SELECT checksum
				FROM public.ariadne_schema_migrations
				WHERE version = $1
			`,
			migration.Name,
		).Scan(&appliedChecksum)
		switch {
		case err == nil:
			if appliedChecksum != migration.Checksum {
				return fmt.Errorf(
					"migration %s checksum changed: applied %s, current %s",
					migration.Name,
					appliedChecksum,
					migration.Checksum,
				)
			}
			continue
		case !errors.Is(err, pgx.ErrNoRows):
			return fmt.Errorf("check migration %s: %w", migration.Name, err)
		}

		transaction, err := connection.Begin(ctx)
		if err != nil {
			return fmt.Errorf("begin migration %s: %w", migration.Name, err)
		}
		if _, err := transaction.Exec(ctx, string(migration.Content)); err != nil {
			_ = transaction.Rollback(ctx)
			return fmt.Errorf("apply migration %s: %w", migration.Name, err)
		}
		if _, err := transaction.Exec(
			ctx,
			`
				INSERT INTO public.ariadne_schema_migrations (version, checksum)
				VALUES ($1, $2)
			`,
			migration.Name,
			migration.Checksum,
		); err != nil {
			_ = transaction.Rollback(ctx)
			return fmt.Errorf("record migration %s: %w", migration.Name, err)
		}
		if err := transaction.Commit(ctx); err != nil {
			return fmt.Errorf("commit migration %s: %w", migration.Name, err)
		}
	}

	return nil
}

func embeddedMigrationManifest() ([]migrationFile, error) {
	entries, err := fs.ReadDir(migrations, "migrations")
	if err != nil {
		return nil, fmt.Errorf("read migrations: %w", err)
	}

	names := make([]string, 0, len(entries))
	for _, entry := range entries {
		if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".up.sql") {
			names = append(names, entry.Name())
		}
	}
	sort.Strings(names)

	manifest := make([]migrationFile, 0, len(names))
	for _, name := range names {
		content, err := migrations.ReadFile("migrations/" + name)
		if err != nil {
			return nil, fmt.Errorf("read migration %s: %w", name, err)
		}
		manifest = append(manifest, migrationFile{
			Name:     name,
			Checksum: fmt.Sprintf("%x", sha256.Sum256(content)),
			Content:  content,
		})
	}
	return manifest, nil
}

func validateSchema(ctx context.Context, database databaseQuerier) error {
	manifest, err := embeddedMigrationManifest()
	if err != nil {
		return err
	}

	rows, err := database.Query(ctx, `
		SELECT version, checksum
		FROM public.ariadne_schema_migrations
		ORDER BY version
	`)
	if err != nil {
		return fmt.Errorf("read applied migrations: %w", err)
	}
	defer rows.Close()

	applied := make(map[string]string, len(manifest))
	for rows.Next() {
		var version, checksum string
		if err := rows.Scan(&version, &checksum); err != nil {
			return fmt.Errorf("scan applied migration: %w", err)
		}
		applied[version] = checksum
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("iterate applied migrations: %w", err)
	}
	if len(applied) != len(manifest) {
		return fmt.Errorf(
			"database migration count = %d, application expects %d",
			len(applied),
			len(manifest),
		)
	}
	for _, migration := range manifest {
		if checksum, exists := applied[migration.Name]; !exists {
			return fmt.Errorf("database migration %s is missing", migration.Name)
		} else if checksum != migration.Checksum {
			return fmt.Errorf("database migration %s checksum does not match", migration.Name)
		}
	}
	return nil
}

func synchronizeAccess(
	ctx context.Context,
	connection *pgxpool.Conn,
	options MigrationOptions,
) error {
	transaction, err := connection.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin access policy transaction: %w", err)
	}
	defer func() {
		rollbackContext, cancel := context.WithTimeout(
			context.Background(),
			5*time.Second,
		)
		defer cancel()
		_ = transaction.Rollback(rollbackContext)
	}()

	if err := validateTenantTableCatalog(ctx, transaction); err != nil {
		return err
	}
	if _, err := transaction.Exec(
		ctx,
		"REVOKE CREATE ON SCHEMA public FROM PUBLIC",
	); err != nil {
		return fmt.Errorf("restrict public schema creation: %w", err)
	}
	if err := synchronizeRuntimeAccess(ctx, transaction, options.RuntimeRole); err != nil {
		return err
	}
	if err := synchronizeIdentityAccess(ctx, transaction, options.IdentityRole); err != nil {
		return err
	}
	if err := validateRuntimePrivileges(ctx, transaction, options.RuntimeRole); err != nil {
		return err
	}
	if err := validateIdentityPrivileges(ctx, transaction, options.IdentityRole); err != nil {
		return err
	}
	if err := transaction.Commit(ctx); err != nil {
		return fmt.Errorf("commit access policies: %w", err)
	}
	return nil
}

func synchronizeRuntimeAccess(ctx context.Context, transaction pgx.Tx, role string) error {
	if err := validateRestrictedRole(ctx, transaction, role, "runtime"); err != nil {
		return err
	}

	roleIdentifier := pgx.Identifier{role}.Sanitize()
	if _, err := transaction.Exec(
		ctx,
		"REVOKE CREATE ON SCHEMA public FROM "+roleIdentifier,
	); err != nil {
		return fmt.Errorf("restrict runtime schema creation: %w", err)
	}
	if _, err := transaction.Exec(
		ctx,
		"GRANT USAGE ON SCHEMA public TO "+roleIdentifier,
	); err != nil {
		return fmt.Errorf("grant runtime schema access: %w", err)
	}

	for _, tableName := range orderedRuntimeTables() {
		permissions, exists := runtimeTablePermissions(tableName)
		if !exists {
			return fmt.Errorf("runtime permissions for %q are not registered", tableName)
		}
		tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
		if _, err := transaction.Exec(
			ctx,
			"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
		); err != nil {
			return fmt.Errorf("reset runtime access to %s: %w", tableName, err)
		}
		if _, err := transaction.Exec(
			ctx,
			"GRANT "+permissions+" ON TABLE "+tableIdentifier+" TO "+roleIdentifier,
		); err != nil {
			return fmt.Errorf("grant runtime access to %s: %w", tableName, err)
		}
	}

	for _, tableName := range []string{"users", "ariadne_schema_migrations"} {
		tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
		if _, err := transaction.Exec(
			ctx,
			"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
		); err != nil {
			return fmt.Errorf("reset runtime access to %s: %w", tableName, err)
		}
	}
	if _, err := transaction.Exec(
		ctx,
		"GRANT SELECT ON TABLE public.ariadne_schema_migrations TO "+roleIdentifier,
	); err != nil {
		return fmt.Errorf("grant runtime migration metadata access: %w", err)
	}
	if _, err := transaction.Exec(
		ctx,
		"GRANT EXECUTE ON FUNCTION public.ariadne_current_organization_id() TO "+
			roleIdentifier,
	); err != nil {
		return fmt.Errorf("grant runtime tenant function access: %w", err)
	}
	return nil
}

func synchronizeIdentityAccess(ctx context.Context, transaction pgx.Tx, role string) error {
	if err := validateRestrictedRole(ctx, transaction, role, "identity"); err != nil {
		return err
	}

	roleIdentifier := pgx.Identifier{role}.Sanitize()
	if _, err := transaction.Exec(
		ctx,
		"REVOKE CREATE ON SCHEMA public FROM "+roleIdentifier,
	); err != nil {
		return fmt.Errorf("restrict identity schema creation: %w", err)
	}
	if _, err := transaction.Exec(
		ctx,
		"GRANT USAGE ON SCHEMA public TO "+roleIdentifier,
	); err != nil {
		return fmt.Errorf("grant identity schema access: %w", err)
	}

	for _, tableName := range orderedRuntimeTables() {
		tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
		if _, err := transaction.Exec(
			ctx,
			"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
		); err != nil {
			return fmt.Errorf("reset identity access to %s: %w", tableName, err)
		}
	}
	for _, tableName := range []string{"users", "ariadne_schema_migrations"} {
		tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
		if _, err := transaction.Exec(
			ctx,
			"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
		); err != nil {
			return fmt.Errorf("reset identity access to %s: %w", tableName, err)
		}
	}

	if _, err := transaction.Exec(
		ctx,
		"GRANT SELECT ON TABLE public.users, public.organizations, "+
			"public.organization_memberships TO "+roleIdentifier,
	); err != nil {
		return fmt.Errorf("grant identity lookup access: %w", err)
	}

	policies := []struct {
		name  string
		table string
	}{
		{name: "organizations_identity_lookup", table: "organizations"},
		{name: "organization_memberships_identity_lookup", table: "organization_memberships"},
	}
	for _, policy := range policies {
		policyIdentifier := pgx.Identifier{policy.name}.Sanitize()
		tableIdentifier := pgx.Identifier{"public", policy.table}.Sanitize()
		if _, err := transaction.Exec(
			ctx,
			"DROP POLICY IF EXISTS "+policyIdentifier+" ON "+tableIdentifier,
		); err != nil {
			return fmt.Errorf("drop identity policy %s: %w", policy.name, err)
		}
		if _, err := transaction.Exec(
			ctx,
			"CREATE POLICY "+policyIdentifier+" ON "+tableIdentifier+
				" FOR SELECT TO "+roleIdentifier+" USING (true)",
		); err != nil {
			return fmt.Errorf("create identity policy %s: %w", policy.name, err)
		}
	}
	return nil
}

func validateRestrictedRole(
	ctx context.Context,
	database databaseQuerier,
	role string,
	purpose string,
) error {
	var (
		exists        bool
		isSuperuser   bool
		canBypassRLS  bool
		canCreateDB   bool
		canCreateRole bool
		canReplicate  bool
		canInherit    bool
	)
	if err := database.QueryRow(ctx, `
		SELECT
			COUNT(*) = 1,
			COALESCE(bool_or(rolsuper), false),
			COALESCE(bool_or(rolbypassrls), false),
			COALESCE(bool_or(rolcreatedb), false),
			COALESCE(bool_or(rolcreaterole), false),
			COALESCE(bool_or(rolreplication), false),
			COALESCE(bool_or(rolinherit), false)
		FROM pg_roles
		WHERE rolname = $1
	`, role).Scan(
		&exists,
		&isSuperuser,
		&canBypassRLS,
		&canCreateDB,
		&canCreateRole,
		&canReplicate,
		&canInherit,
	); err != nil {
		return fmt.Errorf("check %s role: %w", purpose, err)
	}
	if !exists {
		return fmt.Errorf("%s role %q does not exist", purpose, role)
	}
	if isSuperuser ||
		canBypassRLS ||
		canCreateDB ||
		canCreateRole ||
		canReplicate ||
		canInherit {
		return fmt.Errorf(
			"%s role %q must be NOSUPERUSER, NOBYPASSRLS, "+
				"NOCREATEDB, NOCREATEROLE, NOREPLICATION and NOINHERIT",
			purpose,
			role,
		)
	}
	return nil
}