import { ChatOpenAI } from "@langchain/openai";
import { Maxim } from "@maximai/maxim-js";
import { MaximLangchainTracer } from "@maximai/maxim-js/langchain";

// Initialize Maxim
const maxim = new Maxim({
  apiKey: process.env.MAXIM_API_KEY,
});
const logger = await maxim.logger({
  id: process.env.MAXIM_LOG_REPO_ID,
});

if (!logger) {
  throw new Error("logger is not available");
}

// Create the tracer
const maximTracer = new MaximLangchainTracer(logger);

const model = new ChatOpenAI({
  openAIApiKey: process.env.OPENAI_API_KEY,
  modelName: "gpt-4o-mini",
  temperature: 0,
});

const response = await model.invoke("Who is Diego Maradona?", {
  callbacks: [maximTracer],
});

console.log(response.content);

// Clean up resources
await maxim.cleanup();

import { AIMessage, HumanMessage, ToolMessage } from "@langchain/core/messages";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { RunnableLambda } from "@langchain/core/runnables";
import { tool } from "@langchain/core/tools";
import { ChatOpenAI } from "@langchain/openai";
import { Maxim } from "@maximai/maxim-js";
import { MaximLangchainTracer } from "@maximai/maxim-js/langchain";
import { z } from "zod";

// Initialize Maxim
const maxim = new Maxim({
  apiKey: process.env.MAXIM_API_KEY,
});

async function calculatorChainExample() {
  const logger = await maxim.logger({
    id: process.env.MAXIM_LOG_REPO_ID,
  });

  if (!logger) {
    throw new Error("logger is not available");
  }

  const maximTracer = new MaximLangchainTracer(logger);

  // Step 1: Define a calculator tool that can perform basic operations
  const calculatorTool = tool(
    async ({ operation, a, b }) => {
      switch (operation) {
        case "add":
          return a + b;
        case "multiply":
          return a * b;
        case "subtract":
          return a - b;
        case "divide":
          return a / b;
        default:
          throw new Error(`Unknown operation: ${operation}`);
      }
    },
    {
      name: "calculator",
      schema: z.object({
        operation: z.enum(["add", "multiply", "subtract", "divide"]),
        a: z.number(),
        b: z.number(),
      }),
      description: "Performs basic arithmetic operations",
    }
  );

  // Step 2: Create a system prompt for the tool calling chain
  const prompt = ChatPromptTemplate.fromMessages([
    [
      "system",
      "You are a helpful assistant that can perform calculations using tools. Execute tools step by step and provide the final answer.",
    ],
    ["placeholder", "{messages}"],
  ]);

  // Step 3: Create the LLM and bind tools to it
  const llm = new ChatOpenAI({
    openAIApiKey: process.env.OPENAI_API_KEY,
    modelName: "gpt-4o",
    temperature: 0,
    metadata: {
      maxim: {
        generationName: "calculator-model",
        generationTags: {
          testType: "tool-execution",
          complexity: "sequential",
        },
      },
    },
  });

  const llmWithTools = llm.bindTools([calculatorTool]);
  const chain = prompt.pipe(llmWithTools);

  // Step 4: Create a complete tool calling chain that executes tools and gets final response
  const toolChain = RunnableLambda.from(async (userInput: string, config) => {
    // Initialize the conversation with the user's question
    const messages: (HumanMessage | AIMessage | ToolMessage)[] = [new HumanMessage(userInput)];

    // Continue until no more tool calls are needed (max 5 iterations for safety)
    const maxIterations = 5;
    let iteration = 0;

    while (iteration < maxIterations) {
      // Get response from LLM (this may include tool calls)
      const aiMsg = await chain.invoke({ messages }, config);
      messages.push(aiMsg);

      // Check if there are tool calls to execute
      if (!aiMsg.tool_calls || aiMsg.tool_calls.length === 0) {
        // No more tool calls needed - return the final response
        return aiMsg;
      }

      // Execute each tool call and add results to conversation
      for (const toolCall of aiMsg.tool_calls) {
        try {
          // Execute the tool with proper tracing
          const toolResult = await calculatorTool.invoke(toolCall, config);

          // Add tool result to conversation
          if (toolResult instanceof ToolMessage) {
            messages.push(toolResult);
          } else {
            // Create a ToolMessage if needed
            messages.push(new ToolMessage(String(toolResult), toolCall.id || "unknown"));
          }
        } catch (error) {
          // Handle tool execution errors
          const errorMsg = new ToolMessage(`Error executing ${toolCall.name}: ${(error as Error).message}`, toolCall.id || "unknown");
          messages.push(errorMsg);
        }
      }

      iteration++;
    }

    // If we reach max iterations, return the last AI message
    return messages[messages.length - 1];
  }).withConfig({
    // Add metadata for the entire workflow
    metadata: {
      maxim: {
        chainName: "calculator-tool-workflow",
        chainTags: {
          type: "tool-execution",
          workflow: "multi-step-calculation",
          version: "1.0",
        },
      },
    },
  });

  // Step 5: Execute the chain with a complex calculation
  const query = "Calculate 15 * 4 and then add 100 to the result";

  const result = await toolChain.invoke(query, {
    // Add the tracer to capture all operations
    callbacks: [maximTracer],
    // Add comprehensive metadata for this execution
    metadata: {
      maxim: {
        traceName: "Multi-Step Calculation",
        sessionId: "calc_session_001",
        traceTags: {
          category: "mathematics",
          complexity: "multi-step",
          operation_type: "mixed_arithmetic",
          user_type: "demo",
        },
      },
      // Non-Maxim metadata
      query_type: "calculation",
      timestamp: new Date().toISOString(),
    },
  });

  return result;
}

// Execute the example
const calcResult = await calculatorChainExample();
console.log("Calculator Result:", calcResult.content);

// Clean up resources
await maxim.cleanup();