import {
  createProviderToolFactoryWithOutputSchema,
  lazySchema,
  zodSchema,
} from '@ai-sdk/provider-utils';
import { z } from 'zod/v4';

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

export const codeInterpreterOutputSchema = lazySchema(() =>
  zodSchema(
    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 = lazySchema(() =>
  zodSchema(
    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 =
  createProviderToolFactoryWithOutputSchema<
    {
      /**
       * 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',
    inputSchema: codeInterpreterInputSchema,
    outputSchema: codeInterpreterOutputSchema,
  });

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