supabase/apps/studio/components/interfaces/Settings/Database/JitDatabaseAccess/JitDbAccess.utils.ts

import dayjs from 'dayjs'
import { IPv4CidrRange, IPv6CidrRange } from 'ip-num'

import type {
  JitExpiryMode,
  JitIpRangeDraft,
  JitMemberOption,
  JitRoleGrantDraft,
  JitRoleOption,
  JitStatus,
  JitStatusBadge,
  JitUserRule,
  JitUserRuleDraft,
} from './JitDbAccess.types'
import { type DatabaseRolesData, type PgRole } from '@/data/database-roles/database-roles-query'
import type { JitDbAccessMembersData } from '@/data/jit-db-access/jit-db-access-members-query'
import type { OrganizationMembersData } from '@/data/organizations/organization-members-query'
import type { ProjectMembersData } from '@/data/projects/project-members-query'

export function getRelativeDatetimeByMode(mode: JitExpiryMode) {
  if (mode === '1h') return dayjs().add(1, 'hour').toISOString()
  if (mode === '1d') return dayjs().add(1, 'day').toISOString()
  if (mode === '7d') return dayjs().add(7, 'day').toISOString()
  if (mode === '30d') return dayjs().add(30, 'day').toISOString()
  return ''
}

function inferExpiryMode(grant: Pick<JitRoleGrantDraft, 'hasExpiry'>): JitExpiryMode {
  if (!grant.hasExpiry) return 'never'
  return 'custom'
}

export function createEmptyGrant(roleId: string): JitRoleGrantDraft {
  return {
    roleId,
    enabled: false,
    branchesOnly: false,
    expiryMode: '1h',
    hasExpiry: true,
    expiry: getRelativeDatetimeByMode('1h'),
    ipRanges: [createEmptyIpRange()],
  }
}

export function createEmptyIpRange(): JitIpRangeDraft {
  return { value: '' }
}

function parseIpRangeRows(value: JitIpRangeDraft[]) {
  return value.map((item) => item.value.trim()).filter((item) => item.length > 0)
}

function cloneIpRanges(ipRanges: JitIpRangeDraft[]) {
  return ipRanges.map((ipRange) => ({ ...ipRange }))
}

function cloneGrants(grants: JitRoleGrantDraft[]) {
  return grants.map((grant) => ({ ...grant, ipRanges: cloneIpRanges(grant.ipRanges) }))
}

export function createDraft(roleIds: string[]): JitUserRuleDraft {
  return { memberId: '', grants: roleIds.map((roleId) => createEmptyGrant(roleId)) }
}

function mergeRoleIds(baseRoleIds: string[], extraRoleIds: string[]) {
  const seen = new Set<string>()
  const merged: string[] = []

  for (const roleId of [...baseRoleIds, ...extraRoleIds]) {
    if (seen.has(roleId)) continue
    seen.add(roleId)
    merged.push(roleId)
  }

  return merged
}

export function draftFromRule(rule: JitUserRule, baseRoleIds: string[]): JitUserRuleDraft {
  const byRoleId = new Map(rule.grants.map((grant) => [grant.roleId, grant]))
  const mergedRoleIds = mergeRoleIds(
    baseRoleIds,
    rule.grants.map((grant) => grant.roleId)
  )

  return {
    memberId: rule.memberId,
    grants: mergedRoleIds.map((roleId) => {
      const nextGrant = {
        ...createEmptyGrant(roleId),
        ...(byRoleId.get(roleId) ?? {}),
      }

      return {
        ...nextGrant,
        expiryMode: inferExpiryMode(nextGrant),
        ipRanges: cloneIpRanges(nextGrant.ipRanges),
      }
    }),
  }
}

export function computeStatusFromGrants(grants: JitRoleGrantDraft[]): JitStatus {
  const enabledGrants = grants.filter((grant) => grant.enabled)

  let active = 0
  let expired = 0
  let activeIp = 0
  let expiredIp = 0

  enabledGrants.forEach((grant) => {
    const hasIp = parseIpRangeRows(grant.ipRanges).length > 0

    if (!grant.hasExpiry || !grant.expiry) {
      active += 1
      if (hasIp) activeIp += 1
      return
    }

    const isExpired = dayjs(grant.expiry).isValid() && dayjs(grant.expiry).isBefore(dayjs())

    if (isExpired) {
      expired += 1
      if (hasIp) expiredIp += 1
      return
    }

    active += 1
    if (hasIp) activeIp += 1
  })

  return { active, expired, activeIp, expiredIp }
}

function formatBadgeLabel(raw: string, showCount: boolean): string {
  if (showCount) return raw
  // If only one in count:
  // Strip leading "N " and "· N " count segments, then capitalize first letter
  return raw
    .replace(/^\d+\s/, '')
    .replace(/·\s*\d+\s/g, '· ')
    .replace(/^./, (c) => c.toUpperCase())
}

export function getJitStatusDisplay(status: JitStatus): { badges: JitStatusBadge[] } {
  const { active, expired, activeIp } = status
  const badges: JitStatusBadge[] = []
  const showCount = (active > 0 ? 1 : 0) + (expired > 0 ? 1 : 0) > 1

  if (active > 0) {
    const raw = activeIp > 0 ? `${active} active · ${activeIp} IP` : `${active} active`
    badges.push({ label: formatBadgeLabel(raw, showCount), variant: 'success' })
  }

  if (expired > 0) {
    const raw = `${expired} expired`
    badges.push({ label: formatBadgeLabel(raw, showCount), variant: 'default' })
  }

  return { badges }
}

function toUnixSeconds(datetimeIso: string) {
  const value = dayjs(datetimeIso)
  if (!value.isValid()) return undefined
  return value.unix()
}

function isValidCidr(value: string) {
  try {
    if (value.includes(':')) {
      IPv6CidrRange.fromCidr(value)
      return true
    }

    IPv4CidrRange.fromCidr(value)
    return true
  } catch {
    return false
  }
}

export function getInvalidIpRangeRows(value: JitIpRangeDraft[]) {
  return parseIpRangeRows(value).filter((cidr) => !isValidCidr(cidr))
}

function isAssignableJitRole(role: PgRole) {
  return (
    role.canLogin &&
    !role.isSuperuser &&
    !role.name.startsWith('pg_') &&
    (!role.name.startsWith('supabase_') || role.name === 'supabase_read_only_user') &&
    !['pgbouncer', 'authenticator'].includes(role.name)
  )
}

function serializeAllowedNetworks(roleObj: {
  allowed_networks?: {
    allowed_cidrs?: Array<{ cidr: string }>
    allowed_cidrs_v6?: Array<{ cidr: string }>
  }
}) {
  const cidrs = roleObj.allowed_networks?.allowed_cidrs?.map((item) => item.cidr) ?? []
  const cidrsV6 = roleObj.allowed_networks?.allowed_cidrs_v6?.map((item) => item.cidr) ?? []
  return [...cidrs, ...cidrsV6]
}

export function getAssignableJitRoleOptions(
  databaseRoles?: DatabaseRolesData | null
): JitRoleOption[] {
  return (
    databaseRoles
      ?.filter(isAssignableJitRole)
      .map((role) => ({ id: role.name, label: role.name }))
      .sort((a, b) => a.label.localeCompare(b.label)) ?? []
  )
}

export function getJitMemberOptions(
  organizationMembers?: OrganizationMembersData | null,
  projectMembers?: ProjectMembersData | null
): JitMemberOption[] {
  const byId = new Map<string, JitMemberOption>()

  for (const member of organizationMembers ?? []) {
    // JIT rules should only target accepted org members (invites can be expired/pending).
    if (!member.gotrue_id) continue

    const id = member.gotrue_id ?? member.primary_email
    if (!id) continue

    byId.set(id, {
      id,
      email: member.primary_email ?? id,
      name: member.username || undefined,
    })
  }

  for (const member of projectMembers ?? []) {
    const id = member.user_id ?? member.primary_email
    if (!id) continue

    byId.set(id, {
      id,
      email: member.primary_email ?? byId.get(id)?.email ?? id,
      name: member.username ?? byId.get(id)?.name,
    })
  }

  return Array.from(byId.values()).sort((a, b) => a.email.localeCompare(b.email))
}

export function mapJitMembersToUserRules(
  jitMembers: JitDbAccessMembersData | undefined,
  projectMembers: ProjectMembersData | undefined,
  roleOptions: JitRoleOption[]
): JitUserRule[] {
  const memberMap = new Map((projectMembers ?? []).map((member) => [member.user_id, member]))
  const baseRoleIds = roleOptions.map((role) => role.id)

  return (jitMembers ?? []).map((item) => {
    const mappedMember = memberMap.get(item.user_id)
    const assignedRoles: JitRoleGrantDraft[] = (item.user_roles ?? []).map((roleObj) => {
      const roleWithBranchRestriction = roleObj as typeof roleObj & { branches_only?: boolean }
      const expiresAt = typeof roleObj.expires_at === 'number' ? roleObj.expires_at : undefined
      const hasExpiry = typeof expiresAt === 'number'
      const allowedNetworks = serializeAllowedNetworks(roleObj)

      return {
        ...createEmptyGrant(roleObj.role),
        roleId: roleObj.role,
        enabled: true,
        branchesOnly: roleWithBranchRestriction.branches_only ?? false,
        hasExpiry,
        expiryMode: hasExpiry ? 'custom' : 'never',
        expiry: hasExpiry ? new Date(expiresAt * 1000).toISOString() : '',
        ipRanges:
          allowedNetworks.length > 0
            ? allowedNetworks.map((cidr) => ({ value: cidr }))
            : [createEmptyIpRange()],
      }
    })

    const assignedByRoleId = new Map(assignedRoles.map((grant) => [grant.roleId, grant]))
    const allRoleIds = mergeRoleIds(
      baseRoleIds,
      assignedRoles.map((grant) => grant.roleId)
    )
    const grants = allRoleIds.map((roleId) => ({
      ...createEmptyGrant(roleId),
      ...(assignedByRoleId.get(roleId) ?? {}),
      roleId,
    }))

    const email = mappedMember?.primary_email ?? item.user_id
    const name = mappedMember?.username ?? undefined

    return {
      id: item.user_id,
      memberId: item.user_id,
      email,
      name,
      grants: cloneGrants(grants),
      status: computeStatusFromGrants(grants),
    }
  })
}

export function serializeDraftRolesForGrantMutation(draft: JitUserRuleDraft) {
  const serializeAllowedNetworks = (value: JitIpRangeDraft[]) => {
    const cidrs = parseIpRangeRows(value)
    if (cidrs.length === 0) return undefined

    const allowed_cidrs = cidrs.filter((cidr) => !cidr.includes(':')).map((cidr) => ({ cidr }))
    const allowed_cidrs_v6 = cidrs.filter((cidr) => cidr.includes(':')).map((cidr) => ({ cidr }))

    return {
      ...(allowed_cidrs.length > 0 ? { allowed_cidrs } : {}),
      ...(allowed_cidrs_v6.length > 0 ? { allowed_cidrs_v6 } : {}),
    }
  }

  return draft.grants
    .filter((grant) => grant.enabled)
    .map((grant) => {
      const expires_at = grant.hasExpiry ? toUnixSeconds(grant.expiry) : undefined
      const allowed_networks = serializeAllowedNetworks(grant.ipRanges)
      return {
        role: grant.roleId,
        ...(grant.branchesOnly ? { branches_only: true } : {}),
        ...(typeof expires_at === 'number' ? { expires_at } : {}),
        ...(allowed_networks ? { allowed_networks } : {}),
      }
    })
}