import { transformJSONSchema } from '../..//lib/transform-json-schema';
import * as z from 'zod/v4';
import { AnthropicError } from '../../core/error';
import { AutoParseableBetaOutputFormat } from '../../lib/beta-parser';
import { BetaRunnableTool, BetaToolRunContext, Promisable } from '../../lib/tools/BetaRunnableTool';
import { BetaTool, BetaToolResultContentBlockParam } from '../../resources/beta';
/**
 * Creates a JSON schema output format object from the given Zod schema.
 *
 * If this is passed to the `.parse()` method then the response message will contain a
 * `.parsed_output` property that is the result of parsing the content with the given Zod object.
 *
 * This can be passed directly to the `.create()` method but will not
 * result in any automatic parsing, you'll have to parse the response yourself.
 */
export function betaZodOutputFormat<ZodInput extends z.ZodType>(
  zodObject: ZodInput,
): AutoParseableBetaOutputFormat<z.infer<ZodInput>> {
  const jsonSchema = transformJSONSchema(z.toJSONSchema(zodObject, { reused: 'ref' }));

  return {
    type: 'json_schema',
    schema: {
      ...jsonSchema,
    },
    parse: (content) => {
      const output = zodObject.safeParse(JSON.parse(content));

      if (!output.success) {
        throw new AnthropicError(
          `Failed to parse structured output: ${output.error.message} cause: ${output.error.issues}`,
        );
      }

      return output.data;
    },
  };
}

/**
 * Creates a tool using the provided Zod schema that can be passed
 * into the `.toolRunner()` method. The Zod schema will automatically be
 * converted into JSON Schema when passed to the API. The provided function's
 * input arguments will also be validated against the provided schema.
 *
 * Every other field of the tool definition, such as `strict` or `cache_control`,
 * can be passed too and is sent to the API as given.
 */
export function betaZodTool<InputSchema extends z.ZodType>(
  options: {
    name: string;
    inputSchema: InputSchema;
    description: string;
    run: (
      args: z.infer<InputSchema>,
      context?: BetaToolRunContext,
    ) => Promisable<string | Array<BetaToolResultContentBlockParam>>;
    /**
     * Optional cleanup hook for tools that hold process-level resources (e.g. a
     * persistent shell). `client.beta.sessions.events.toolRunner` calls it once
     * when iteration ends.
     */
    close?: () => void | Promise<void>;
  } & Omit<BetaTool, 'type' | 'name' | 'description' | 'input_schema'>,
): BetaRunnableTool<z.infer<InputSchema>> {
  const { name, inputSchema, description, run, ...extraProps } = options;
  const jsonSchema = z.toJSONSchema(inputSchema, { reused: 'ref' });

  if (jsonSchema.type !== 'object') {
    throw new Error(`Zod schema for tool "${name}" must be an object, but got ${jsonSchema.type}`);
  }

  // TypeScript doesn't narrow the type after the runtime check, so we need to assert it
  const objectSchema = jsonSchema as typeof jsonSchema & { type: 'object' };

  return {
    ...extraProps,
    type: 'custom',
    name,
    input_schema: objectSchema,
    description,
    run,
    parse: (args: unknown) => inputSchema.parse(args) as z.infer<InputSchema>,
  };
}