import { createProviderDefinedToolFactoryWithOutputSchema } from "@ai-sdk/provider-utils"
import { z } from "zod/v4"

export const codeInterpreterInputSchema = z.object({
  code: z.string().nullish(),
  containerId: z.string(),
})

export const codeInterpreterOutputSchema = z.object({
  outputs: z
    .array(
      z.discriminatedUnion("type", [
        z.object({ type: z.literal("logs"), logs: z.string() }),
        z.object({ type: z.literal("image"), url: z.string() }),
      ]),
    )
    .nullish(),
})

export const codeInterpreterArgsSchema = z.object({
  container: z
    .union([
      z.string(),
      z.object({
        fileIds: z.array(z.string()).optional(),
      }),
    ])
    .optional(),
})

type CodeInterpreterArgs = {
  /**
   * The code interpreter container.
   * Can be a container ID
   * or an object that specifies uploaded file IDs to make available to your code.
   */
  container?: string | { fileIds?: string[] }
}

export const codeInterpreterToolFactory = createProviderDefinedToolFactoryWithOutputSchema<
  {
    /**
     * The code to run, or null if not available.
     */
    code?: string | null

    /**
     * The ID of the container used to run the code.
     */
    containerId: string
  },
  {
    /**
     * The outputs generated by the code interpreter, such as logs or images.
     * Can be null if no outputs are available.
     */
    outputs?: Array<
      | {
          type: "logs"

          /**
           * The logs output from the code interpreter.
           */
          logs: string
        }
      | {
          type: "image"

          /**
           * The URL of the image output from the code interpreter.
           */
          url: string
        }
    > | null
  },
  CodeInterpreterArgs
>({
  id: "openai.code_interpreter",
  name: "code_interpreter",
  inputSchema: codeInterpreterInputSchema,
  outputSchema: codeInterpreterOutputSchema,
})

export const codeInterpreter = (
  args: CodeInterpreterArgs = {}, // default
) => {
  return codeInterpreterToolFactory(args)
}