package database

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"

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

const (
	permissionRead       = "SELECT"
	permissionAppend     = "SELECT, INSERT"
	permissionMutate     = "SELECT, INSERT, UPDATE"
	permissionFullAccess = "SELECT, INSERT, UPDATE, DELETE"
)

// runtimeTableAccess is intentionally explicit. Adding an organization-owned
// table without registering its least-privilege access must fail migration and
// server startup instead of silently exposing it.
var runtimeTableAccess = map[string]string{
	"organizations":                             permissionRead,
	"organization_memberships":                  permissionMutate,
	"projects":                                  permissionMutate,
	"ingestion_sources":                         permissionMutate,
	"repositories":                              permissionMutate,
	"runtime_environments":                      permissionMutate,
	"analysis_target_rules":                     permissionMutate,
	"analysis_target_rule_runtime_environments": permissionFullAccess,
	"error_groups":                              permissionMutate,
	"error_occurrences":                         permissionAppend,
	"remediation_runs":                          permissionMutate,
	"remediation_executions":                    permissionMutate,
	"remediation_run_runtime_targets":           permissionAppend,
	"remediation_approvals":                     permissionAppend,
	"merge_requests":                            permissionMutate,
	"audit_events":                              permissionAppend,
}

type databaseQuerier interface {
	Query(context.Context, string, ...any) (pgx.Rows, error)
	QueryRow(context.Context, string, ...any) pgx.Row
}

func orderedRuntimeTables() []string {
	tables := make([]string, 0, len(runtimeTableAccess))
	for table := range runtimeTableAccess {
		tables = append(tables, table)
	}
	sort.Strings(tables)
	return tables
}

func validateTenantTableCatalog(ctx context.Context, database databaseQuerier) error {
	rows, err := database.Query(ctx, `
		SELECT
			table_class.relname,
			table_class.relrowsecurity,
			table_class.relforcerowsecurity,
			COUNT(policy.oid),
			COUNT(policy.oid) FILTER (
				WHERE 0::oid = ANY(policy.polroles)
			),
			COALESCE(
				BOOL_OR(
					0::oid = ANY(policy.polroles)
					AND policy.polcmd = '*'
					AND pg_get_expr(
						policy.polqual,
						policy.polrelid
					) LIKE '%ariadne_current_organization_id()%'
					AND pg_get_expr(
						policy.polwithcheck,
						policy.polrelid
					) LIKE '%ariadne_current_organization_id()%'
				),
				false
			)
		FROM pg_class table_class
		JOIN pg_namespace namespace
		  ON namespace.oid = table_class.relnamespace
		LEFT JOIN pg_policy policy
		  ON policy.polrelid = table_class.oid
		WHERE namespace.nspname = 'public'
		  AND table_class.relkind IN ('r', 'p')
		  AND (
		      table_class.relname = 'organizations'
		      OR EXISTS (
		          SELECT 1
		          FROM information_schema.columns column_info
		          WHERE column_info.table_schema = 'public'
		            AND column_info.table_name = table_class.relname
		            AND column_info.column_name = 'organization_id'
		      )
		  )
		GROUP BY
			table_class.relname,
			table_class.relrowsecurity,
			table_class.relforcerowsecurity
		ORDER BY table_class.relname
	`)
	if err != nil {
		return fmt.Errorf("list tenant table security: %w", err)
	}
	defer rows.Close()

	discovered := make(map[string]struct{}, len(runtimeTableAccess))
	for rows.Next() {
		var (
			tableName             string
			rlsEnabled            bool
			rlsForced             bool
			policyCount           int64
			publicPolicyCount     int64
			hasScopedPublicPolicy bool
		)
		if err := rows.Scan(
			&tableName,
			&rlsEnabled,
			&rlsForced,
			&policyCount,
			&publicPolicyCount,
			&hasScopedPublicPolicy,
		); err != nil {
			return fmt.Errorf("scan tenant table security: %w", err)
		}
		if _, registered := runtimeTableAccess[tableName]; !registered {
			return fmt.Errorf("tenant table %q has no explicit runtime access policy", tableName)
		}
		if !rlsEnabled || !rlsForced || policyCount == 0 {
			return fmt.Errorf(
				"tenant table %q must have enabled and forced RLS with at least one policy",
				tableName,
			)
		}
		if publicPolicyCount != 1 || !hasScopedPublicPolicy {
			return fmt.Errorf(
				"tenant table %q must have exactly one organization-scoped public policy",
				tableName,
			)
		}
		discovered[tableName] = struct{}{}
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("iterate tenant table security: %w", err)
	}

	for tableName := range runtimeTableAccess {
		if _, exists := discovered[tableName]; !exists {
			return fmt.Errorf("registered tenant table %q does not exist", tableName)
		}
	}
	return nil
}

func validateRuntimePrivileges(
	ctx context.Context,
	database databaseQuerier,
	roleName string,
) error {
	for _, tableName := range orderedRuntimeTables() {
		if err := validateTablePrivileges(
			ctx,
			database,
			roleName,
			"runtime",
			tableName,
			runtimeTableAccess[tableName],
		); err != nil {
			return err
		}
	}

	for tableName, expected := range map[string]string{
		"users":                     "",
		"ariadne_schema_migrations": permissionRead,
	} {
		if err := validateTablePrivileges(
			ctx,
			database,
			roleName,
			"runtime",
			tableName,
			expected,
		); err != nil {
			return err
		}
	}

	var canResolveTenant bool
	if err := database.QueryRow(
		ctx,
		`
			SELECT has_function_privilege(
				$1,
				'public.ariadne_current_organization_id()',
				'EXECUTE'
			)
		`,
		roleName,
	).Scan(&canResolveTenant); err != nil {
		return fmt.Errorf("check runtime tenant function privilege: %w", err)
	}
	if !canResolveTenant {
		return errors.New("runtime role must execute tenant scope function")
	}
	return nil
}

func validateIdentityPrivileges(
	ctx context.Context,
	database databaseQuerier,
	roleName string,
) error {
	for _, tableName := range orderedRuntimeTables() {
		expected := ""
		if tableName == "organizations" ||
			tableName == "organization_memberships" {
			expected = permissionRead
		}
		if err := validateTablePrivileges(
			ctx,
			database,
			roleName,
			"identity",
			tableName,
			expected,
		); err != nil {
			return err
		}
	}

	for tableName, expected := range map[string]string{
		"users":                     permissionRead,
		"ariadne_schema_migrations": "",
	} {
		if err := validateTablePrivileges(
			ctx,
			database,
			roleName,
			"identity",
			tableName,
			expected,
		); err != nil {
			return err
		}
	}

	var canResolveTenant bool
	if err := database.QueryRow(
		ctx,
		`
			SELECT has_function_privilege(
				$1,
				'public.ariadne_current_organization_id()',
				'EXECUTE'
			)
		`,
		roleName,
	).Scan(&canResolveTenant); err != nil {
		return fmt.Errorf("check identity tenant function privilege: %w", err)
	}
	if canResolveTenant {
		return errors.New("identity role must not execute tenant scope function")
	}
	return nil
}

func validateTablePrivileges(
	ctx context.Context,
	database databaseQuerier,
	roleName string,
	purpose string,
	tableName string,
	expected string,
) error {
	for _, permission := range []string{"SELECT", "INSERT", "UPDATE", "DELETE"} {
		var granted bool
		if err := database.QueryRow(
			ctx,
			"SELECT has_table_privilege($1, $2, $3)",
			roleName,
			"public."+tableName,
			permission,
		).Scan(&granted); err != nil {
			return fmt.Errorf(
				"check %s %s privilege on %s: %w",
				purpose,
				permission,
				tableName,
				err,
			)
		}
		shouldBeGranted := strings.Contains(expected, permission)
		if granted != shouldBeGranted {
			return fmt.Errorf(
				"%s %s privilege on %s = %t, want %t",
				purpose,
				permission,
				tableName,
				granted,
				shouldBeGranted,
			)
		}
	}
	return nil
}

func runtimeTablePermissions(tableName string) (string, bool) {
	permissions, exists := runtimeTableAccess[tableName]
	return permissions, exists
}