import {
  AUTOMATION_IMAGE_GENERATION_ENTRYPOINT_PREFIX,
  DEFAULT_GATEWAY_IMAGE_MODEL_ID,
  imageGenerationCompletionSchema,
  imageGenerationInputSchema,
  imageGenerationResultSchema,
  platformImageGenerationInputSchema,
  platformImageGenerationStartSchema,
  type GatewayImageModelId,
  type ImageGenerationResult as BaseImageGenerationResult,
} from "@automate.ax/api-contract/runtime"
import { automationInvokedTriggerDefinition } from "@automate.ax/catalog/triggers/core-invocation"
import { createDeepInfra } from "@ai-sdk/deepinfra"
import { createFireworks } from "@ai-sdk/fireworks"
import { createGoogleGenerativeAI } from "@ai-sdk/google"
import { createOpenAI } from "@ai-sdk/openai"
import { createTogetherAI } from "@ai-sdk/togetherai"
import { createXai } from "@ai-sdk/xai"
import {
  createGateway,
  generateImage as generateAiImage,
  type ImageModel,
} from "ai"
import * as z from "zod"
import { defineAction, type ActionObjectInput } from "../../automation/actions"
import type {
  IntegrationAccountReference,
  ResolvedIntegrationAccount,
} from "../../automation/integrations"
import { serializedIntegrationAccountDefinition } from "../../automation/integrations"
import { getCurrentHookScopePath } from "../../automation/runtime"
import {
  correlate,
  getCurrentContextPrerequisite,
  getCurrentSignalPrerequisites,
  group,
  scope,
  withoutSignalPrerequisites,
  withContextPrerequisite,
  withPrerequisites,
} from "../../automation/signal-operators"
import {
  FailedSignalError,
  transform,
  type Signal,
} from "../../automation/signal-protocol"
import { createSubscription } from "../../automation/subscription"

const AI_IMAGE_ACCOUNT_SERVICE_IDS = [
  "deepinfra",
  "fireworks",
  "google-genai",
  "openai",
  "togetherai",
  "vercel-ai-gateway",
  "xai",
] as const

type AiImageAccountServiceId = (typeof AI_IMAGE_ACCOUNT_SERVICE_IDS)[number]

/** Serializable AI SDK image-generation result. */
export type GenerateImageResult = BaseImageGenerationResult

interface AiImageModelIdsByService {
  deepinfra: Parameters<ReturnType<typeof createDeepInfra>["image"]>[0]
  fireworks: Parameters<ReturnType<typeof createFireworks>["image"]>[0]
  "google-genai": Parameters<
    ReturnType<typeof createGoogleGenerativeAI>["image"]
  >[0]
  openai: Parameters<ReturnType<typeof createOpenAI>["image"]>[0]
  togetherai: Parameters<ReturnType<typeof createTogetherAI>["image"]>[0]
  "vercel-ai-gateway": GatewayImageModelId
  xai: Parameters<ReturnType<typeof createXai>["image"]>[0]
}

/** Image model IDs suggested for one provider, including newly released IDs. */
export type GenerateImageModelId<
  TServiceId extends AiImageAccountServiceId = AiImageAccountServiceId,
> = AiImageModelIdsByService[TServiceId] | (string & {})

const DEFAULT_AI_IMAGE_MODELS = {
  deepinfra: "black-forest-labs/FLUX-1-schnell",
  fireworks: "accounts/fireworks/models/flux-1-schnell-fp8",
  "google-genai": "imagen-4.0-fast-generate-001",
  openai: "gpt-image-1-mini",
  togetherai: "black-forest-labs/FLUX.1-schnell-Free",
  "vercel-ai-gateway": DEFAULT_GATEWAY_IMAGE_MODEL_ID,
  xai: "grok-imagine-image",
} as const satisfies Record<AiImageAccountServiceId, GenerateImageModelId>

const openAiImageFetch = Object.assign(fetchOpenAiImage, {
  preconnect: fetch.preconnect,
})

const AI_IMAGE_MODEL_FACTORIES = {
  deepinfra: (apiKey: string, model: GenerateImageModelId) =>
    createDeepInfra({ apiKey }).image(model),
  fireworks: (apiKey: string, model: GenerateImageModelId) =>
    createFireworks({ apiKey }).image(model),
  "google-genai": (apiKey: string, model: GenerateImageModelId) =>
    createGoogleGenerativeAI({ apiKey }).image(model),
  openai: (apiKey: string, model: GenerateImageModelId) =>
    createOpenAI({ apiKey, fetch: openAiImageFetch }).image(model),
  togetherai: (apiKey: string, model: GenerateImageModelId) =>
    createTogetherAI({ apiKey }).image(model),
  "vercel-ai-gateway": (apiKey: string, model: GenerateImageModelId) =>
    createGateway({ apiKey }).image(model),
  xai: (apiKey: string, model: GenerateImageModelId) =>
    createXai({ apiKey }).image(model),
} satisfies Record<
  AiImageAccountServiceId,
  (apiKey: string, model: GenerateImageModelId) => ImageModel
>

const generateProviderImageAction = defineAction("Generate image")
  .account(AI_IMAGE_ACCOUNT_SERVICE_IDS)
  .input(imageGenerationInputSchema)
  .output(imageGenerationResultSchema)
  .retry({ replaySafety: "unsafe" })
  .handler(async ({ account, input }) =>
    normalizeImageGenerationResult(
      await generateAiImage({
        ...input,
        model: createProviderImageModel(account, input.model),
        prompt: await normalizeImagePrompt(input.prompt),
      }),
    ),
  )

const startPlatformImageGenerationAction = defineAction("Generate image")
  .input(platformImageGenerationInputSchema)
  .output(platformImageGenerationStartSchema)
  .retry({ replaySafety: "safe" })
  .handler(
    async ({ input, runtime }) => await runtime.startImageGeneration(input),
  )

/** Inputs accepted by the built-in image-generation action. */
export type GenerateImageOptions<
  TServiceId extends AiImageAccountServiceId = AiImageAccountServiceId,
> = Omit<ActionObjectInput<typeof imageGenerationInputSchema>, "model"> & {
  /** User-managed provider account. Omit to use the platform AI Gateway. */
  account?: IntegrationAccountReference<TServiceId>

  /** Provider image model ID. Defaults according to the selected account. */
  model?:
    | GenerateImageModelId<NoInfer<TServiceId>>
    | Exclude<
        ActionObjectInput<typeof imageGenerationInputSchema>["model"],
        GatewayImageModelId | undefined
      >
}

/**
 * Generates or edits images through the platform AI Gateway.
 *
 * @param options - Prompt, source images, model, and generation settings.
 */
export function generateImage(
  options: Omit<GenerateImageOptions<"vercel-ai-gateway">, "account"> & {
    account?: never
  },
): Signal<GenerateImageResult>

/**
 * Generates or edits images through an explicit provider account.
 *
 * @param options - Provider account, prompt, source images, and settings.
 */
export function generateImage<const TServiceId extends AiImageAccountServiceId>(
  options: GenerateImageOptions<TServiceId> & {
    account: IntegrationAccountReference<TServiceId>
  },
): Signal<GenerateImageResult>

/**
 * Generates or edits images when the account may vary at authoring time.
 *
 * @param options - Account, prompt, source images, and settings.
 */
export function generateImage(
  options: GenerateImageOptions,
): Signal<GenerateImageResult>

export function generateImage(options: GenerateImageOptions): unknown {
  const { account, model, ...settings } = options
  if (!account) {
    return generatePlatformImage({
      ...settings,
      model: model ?? DEFAULT_GATEWAY_IMAGE_MODEL_ID,
    })
  }

  return generateProviderImageAction(
    {
      ...settings,
      model:
        model ??
        DEFAULT_AI_IMAGE_MODELS[
          account[serializedIntegrationAccountDefinition].serviceId
        ],
    },
    { account },
  )
}

/**
 * Builds the internal invocation and correlation flow for Platform AI.
 *
 * @param input - Platform image settings and compatible signals.
 */
function generatePlatformImage(
  input: ActionObjectInput<typeof imageGenerationInputSchema>,
) {
  const prerequisites = getCurrentSignalPrerequisites()
  const contextPrerequisite = getCurrentContextPrerequisite()
  return group({ name: "Generate image", presentation: "hidden" }, () =>
    withoutSignalPrerequisites(() => {
      const entrypoint = `${AUTOMATION_IMAGE_GENERATION_ENTRYPOINT_PREFIX}${getCurrentHookScopePath().join(".")}`
      const completed = createSubscription<
        z.output<typeof imageGenerationCompletionSchema>
      >(automationInvokedTriggerDefinition, { entrypoint }, undefined, {
        inferActionBoundary: false,
      })
      const start = () =>
        startPlatformImageGenerationAction({
          ...input,
          entrypoint,
        })
      const startWithContext = () =>
        contextPrerequisite
          ? withContextPrerequisite(contextPrerequisite, start)
          : start()
      return transform(
        [
          correlate([
            (prerequisites
              ? scope(() => withPrerequisites(prerequisites, startWithContext))
              : startWithContext()
            ).keyBy(({ key }) => key),
            completed.keyBy(({ key }) => key),
          ]),
          completed,
        ],
        (_matched, { outcome }) => {
          if (outcome.status === "failed") {
            throw new FailedSignalError(outcome.failure)
          }
          return outcome.result
        },
      )
    }),
  )
}

/**
 * Resolves one user-managed image model from its API-key account.
 *
 * @param account - Resolved provider account.
 * @param model - Provider-native image model ID.
 */
function createProviderImageModel(
  account: ResolvedIntegrationAccount<AiImageAccountServiceId>,
  model: GenerateImageModelId,
): ImageModel {
  const { apiKey } = z
    .object({ apiKey: z.string().min(1) })
    .parse(account.secret)

  return AI_IMAGE_MODEL_FACTORIES[account.serviceId](apiKey, model)
}

/**
 * Selects the serializable AI SDK image result fields exposed by the action.
 *
 * @param result - Completed AI SDK image-generation result.
 */
function normalizeImageGenerationResult(
  result: Awaited<ReturnType<typeof generateAiImage>>,
): GenerateImageResult {
  return imageGenerationResultSchema.parse({
    images: result.images.map(
      (image) =>
        new Blob([new Uint8Array(image.uint8Array)], {
          type: image.mediaType,
        }),
    ),
    providerMetadata: imageGenerationResultSchema.shape.providerMetadata.parse(
      JSON.parse(JSON.stringify(result.providerMetadata)),
    ),
    responses: result.responses.map((response) => ({
      headers: response.headers,
      modelId: response.modelId,
      timestamp: response.timestamp.toISOString(),
    })),
    usage: result.usage,
    warnings: result.warnings,
  })
}

/**
 * Converts canonical Blob prompt content to the AI SDK's data URL input.
 *
 * @param prompt - Text or structured image prompt.
 */
async function normalizeImagePrompt(
  prompt: z.output<typeof imageGenerationInputSchema>["prompt"],
) {
  if (typeof prompt === "string") return prompt
  return {
    ...prompt,
    images: await Promise.all(prompt.images.map(normalizeImageDataContent)),
    mask:
      prompt.mask === undefined
        ? undefined
        : await normalizeImageDataContent(prompt.mask),
  }
}

/**
 * Converts one Blob prompt input while preserving its media type.
 *
 * @param content - One AI SDK-compatible image input.
 */
async function normalizeImageDataContent(
  content: string | Uint8Array | ArrayBuffer | Blob,
) {
  if (!(content instanceof Blob)) return content
  const bytes = new Uint8Array(await content.arrayBuffer())
  return content.type
    ? `data:${content.type};base64,${Buffer.from(bytes).toString("base64")}`
    : bytes
}

/**
 * Adds filenames to OpenAI multipart image uploads for Bun compatibility.
 *
 * @param input - Request target passed by the OpenAI provider.
 * @param init - Request options passed by the OpenAI provider.
 */
function fetchOpenAiImage(
  input: Parameters<typeof fetch>[0],
  init?: Parameters<typeof fetch>[1],
) {
  if (!(init?.body instanceof FormData)) return fetch(input, init)

  const formData = new FormData()
  let blobIndex = 0
  for (const [name, value] of init.body.entries()) {
    const entry: unknown = value
    if (!(entry instanceof Blob)) {
      formData.append(name, value)
      continue
    }

    blobIndex += 1
    const mediaSubtype = entry.type.split("/")[1]?.split(";")[0]
    formData.append(
      name,
      entry,
      `image-${blobIndex}.${mediaSubtype === "jpeg" ? "jpg" : mediaSubtype || "bin"}`,
    )
  }

  return fetch(input, { ...init, body: formData })
}