import { z } from "zod";

type ToolContext = {
  sessionID: string;
};

type ToolDefinition<Shape extends z.ZodRawShape> = {
  description: string;
  args: Shape;
  execute: (
    args: z.infer<z.ZodObject<Shape>>,
    context: ToolContext,
  ) => Promise<string> | string;
};

const defineTool = <Shape extends z.ZodRawShape>(
  definition: ToolDefinition<Shape>,
) => {
  const inputSchema = z.object(definition.args);
  return {
    description: definition.description,
    input: z.toJSONSchema(inputSchema),
    async execute(input: unknown, context: ToolContext) {
      return {
        content: await definition.execute(inputSchema.parse(input), context),
      };
    },
  };
};

export const tool = Object.assign(defineTool, { schema: z });