import type { AutomateClient } from "@automate.ax/client"
import type { ActionRuntime } from "./actions"

interface RuntimeAutomateClients {
  forApiKey: (apiKey: string) => AutomateClient
  runtime: AutomateClient
}

const AUTOMATE_CLIENTS = new WeakMap<ActionRuntime, RuntimeAutomateClients>()

/**
 * Associates one action runtime with its execution-authenticated client.
 *
 * @param runtime - Action runtime receiving the client.
 * @param client - Client authenticated by the active runtime token.
 * @param forApiKey - Creates a client on the same runtime transport for an
 *   explicit connected-account credential.
 */
export function setRuntimeAutomateClient(
  runtime: ActionRuntime,
  client: AutomateClient,
  forApiKey: (apiKey: string) => AutomateClient,
) {
  AUTOMATE_CLIENTS.set(runtime, { forApiKey, runtime: client })
}

/**
 * Returns the execution-authenticated client owned by an action runtime.
 *
 * @param runtime - Action runtime whose client is required.
 * @throws When the runtime was not created with execution authentication.
 */
export function getRuntimeAutomateClient(runtime: ActionRuntime) {
  const clients = AUTOMATE_CLIENTS.get(runtime)
  if (!clients) {
    throw new Error("Automate.ax runtime authentication is unavailable.")
  }
  return clients.runtime
}

/**
 * Creates a connected-account client on the active runtime API transport.
 *
 * @param runtime - Action runtime whose API transport is required.
 * @param apiKey - Connected Automate.ax organization API key.
 * @throws When the runtime was not created with management API access.
 */
export function getRuntimeAutomateClientForApiKey(
  runtime: ActionRuntime,
  apiKey: string,
) {
  const clients = AUTOMATE_CLIENTS.get(runtime)
  if (!clients) {
    throw new Error("Automate.ax runtime authentication is unavailable.")
  }
  return clients.forApiKey(apiKey)
}