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 { 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>> {
  let jsonSchema = z.toJSONSchema(zodObject, { reused: 'ref' });

  jsonSchema = transformJSONSchema(jsonSchema);

  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.
 */
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>;
}): BetaRunnableTool<z.infer<InputSchema>> {
  const jsonSchema = z.toJSONSchema(options.inputSchema, { reused: 'ref' });

  if (jsonSchema.type !== 'object') {
    throw new Error(`Zod schema for tool "${options.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 {
    type: 'custom',
    name: options.name,
    input_schema: objectSchema,
    description: options.description,
    run: options.run,
    parse: (args: unknown) => options.inputSchema.parse(args) as z.infer<InputSchema>,
    ...(options.close ? { close: options.close } : {}),
  };
}