/**
 * JWT access-token helpers for Databricks Apps.
 *
 * Parses bearer tokens and `x-forwarded-access-token` payloads to read
 * OAuth scopes and other claims from incoming HTTP requests.
 *
 * @module
 */

import { forEachHeaderValue, type HeaderLike } from "./http.ts";
import * as json from "./json.ts";
import { isLevelEnabled, logger } from "./log.ts";

const BEARER_PREFIX_REGEX = /^bearer\s+/i;
const SPLIT_REGEX = /\s+|\s*,\s*/;
const log = logger("shared/token");

/**
 * Header the Databricks Apps front door forwards the caller's OAuth (OBO) token
 * on - the default source for every helper here.
 *
 * Exported because it is a wire contract several packages branch on, not just a
 * default argument: `@dbx-tools/appkit`'s `identity` module decides whether a
 * request can use OBO auth at all by its presence, and `@dbx-tools/tunnel`
 * must know the name to strip inbound copies of it. One spelling, one place.
 */
export const ACCESS_TOKEN_HEADER = "x-forwarded-access-token";

/** Header the front door forwards the caller's user id on. */
export const USER_ID_HEADER = "x-forwarded-user";

/** Header the front door forwards the caller's email on. */
export const USER_EMAIL_HEADER = "x-forwarded-email";

/**
 * Decode a JWT segment (base64url with standard base64 padding), or
 * `undefined` when the segment is not valid base64.
 */
function decodeJwtSegment(segment: string): string | undefined {
  const normalized = segment.replace(/-/g, "+").replace(/_/g, "/");
  const padded = normalized + "=".repeat((4 - (normalized.length % 4)) % 4);
  try {
    const binary = atob(padded);
    const bytes = Uint8Array.from(binary, (ch) => ch.charCodeAt(0));
    return new TextDecoder().decode(bytes);
  } catch {
    return undefined;
  }
}

/**
 * Decode the JWT payload from a bearer token string or request headers.
 *
 * When `input` is header-like, walks every value for `headerName` (default
 * `x-forwarded-access-token`) and merges claim keys left-to-right.
 */
export function getAccessTokenPayload(
  input: HeaderLike | string,
  headerName: string | undefined = ACCESS_TOKEN_HEADER,
): Record<string, unknown> {
  let accessTokenPayload: Record<string, unknown> | undefined;
  if (!(typeof input === "string")) {
    if (headerName) {
      forEachHeaderValue(input, headerName, (value: string) => {
        for (const [payloadKey, payloadValue] of Object.entries(getAccessTokenPayload(value))) {
          if (!accessTokenPayload || !(payloadKey in accessTokenPayload)) {
            if (!accessTokenPayload) accessTokenPayload = {};
            accessTokenPayload[payloadKey] = payloadValue;
          }
        }
      });
    }
  } else {
    input = input.trim();
    if (input) {
      const match = BEARER_PREFIX_REGEX.exec(input);
      if (match) {
        const endIndex = match.index + match[0].length;
        input = input.slice(endIndex);
      }
      const parts = input.split(".", 4);
      if (parts.length === 2 || parts.length === 3) {
        // A malformed segment yields undefined; fall through to empty object.
        accessTokenPayload = json.parseRecord(decodeJwtSegment(parts[1]!));
      }
    }
  }
  if (!accessTokenPayload) accessTokenPayload = {};
  log.debug("getAccessTokenPayload", { accessTokenPayload });
  return accessTokenPayload;
}

/**
 * OAuth scopes from the access-token JWT on `input`.
 *
 * Defaults to the `x-forwarded-access-token` header, which Databricks Apps
 * forward from the browser session.
 */
export function getAccessTokenScopes(
  input: HeaderLike | string,
  headerName?: string,
): Iterable<string> {
  const payload = getAccessTokenPayload(input, headerName);
  let accessTokenScopes: Iterable<string> = iterateClaims(payload.scope);
  if (isLevelEnabled("debug")) {
    accessTokenScopes = Array.from(accessTokenScopes);
    log.debug("getAccessTokenScopes", { accessTokenScopes });
  }
  return accessTokenScopes;
}

/**
 * Return true when `scopes` contains any value from `allowed`.
 *
 * Accepts the raw {@link MASTRA_SCOPES_KEY} array stamped on
 * `RequestContext` as well as iterables from {@link getAccessTokenScopes}.
 */
export function includesAccessTokenScope(scopes: unknown, allowed: readonly string[]): boolean {
  if (allowed.length) {
    for (const scope of iterateClaims(scopes)) {
      if (allowed.includes(scope)) {
        log.debug("includesAccessTokenScope", { scope, allowed });
        return true;
      }
    }
  }
  return false;
}

/** Split a JWT `scope` claim string or array into individual scope tokens. */
function* iterateClaims(input: unknown, distinct = true): Iterable<string> {
  const seen = distinct ? new Set<string>() : undefined;

  function* visit(value: unknown): Iterable<string> {
    if (typeof value === "string") {
      for (const claim of value.split(SPLIT_REGEX)) {
        const token = claim.trim();
        if (!token || seen?.has(token)) {
          continue;
        }
        seen?.add(token);
        yield token;
      }
    } else if (Array.isArray(value)) {
      for (const item of value) {
        yield* visit(item);
      }
    }
  }

  yield* visit(input);
}