import { icepick } from "@hatchet-dev/icepick";
import { myTool, anotherTool, calculatorTool } from "@tools";

export const myToolbox = icepick.toolbox({
  tools: [
    myTool,
    anotherTool,
    calculatorTool,
  ],
});

export const myAgent = icepick.agent({
  name: "my-agent",
  inputSchema: z.object({ message: z.string() }),
  outputSchema: z.object({ response: z.string() }),
  fn: async (input, ctx) => {
    const result = await myToolbox.pickAndRun({
      prompt: input.message,
    });

    switch (result.name) {
      case "myTool":
        return { response: `Tool result: ${result.output}` };
      case "anotherTool":
        return { response: `Another result: ${result.output}` };
      case "calculatorTool":
        return { response: `Calculation: ${result.output}` };
      default:
        return myToolbox.assertExhaustive(result);
    }
  },
});

const selections = await myToolbox.pick({
  prompt: "Calculate the square root of 16",
});

// selections contains the selected tool names and generated inputs
console.log(selections); // [{ name: "calculatorTool", input: { operation: "sqrt", value: 16 } }]

const results = await myToolbox.pickAndRun({
  prompt: "Process the data and then send a notification",
  maxTools: 2,
});

// results is an array when maxTools > 1
results.forEach(result => {
  console.log(`${result.name}: ${result.output}`);
});

export const smartAgent = icepick.agent({
  name: "smart-agent",
  inputSchema: z.object({ task: z.string() }),
  outputSchema: z.object({ result: z.string() }),
  fn: async (input, ctx) => {
    // The AI selects the appropriate tool for the task
    const result = await myToolbox.pickAndRun({
      prompt: `Help me with this task: ${input.task}`,
    });

    return {
      result: `Completed using ${result.name}: ${result.output}`,
    };
  },
});