import { tool, createAgent, AgentState } from "langchain";
import { wrapModelCall, ModelRequest, ModelResponse } from "langchain/agents/middleware";
import { Command } from "@langchain/langgraph";
import { ToolMessage } from "langchain/messages";
import * as z from "zod";

// 1. Define state with current_step tracker
const SupportState = z.object({  
  currentStep: z.string().default("triage"),  
  warrantyStatus: z.string().optional()
});

// 2. Tools update currentStep via Command
const recordWarrantyStatus = tool(
  async ({ status }, config) => {
    return new Command({  
      update: {  
        messages: [  
          new ToolMessage({
            content: `Warranty status recorded: ${status}`,
            tool_call_id: config.toolCall?.id!
          })
        ],
        warrantyStatus: status,
        // Transition to next step
        currentStep: "specialist",  
      }
    });
  },
  {
    name: "record_warranty_status",
    description: "Record warranty status and transition",
    schema: z.object({ status: z.string() })
  }
);

// 3. Middleware applies dynamic configuration based on currentStep
const applyStepConfig = wrapModelCall(  
  async (request: ModelRequest, handler) => {
    const step = request.state.currentStep || "triage";  

    // Map steps to their configurations
    const configs = {
      triage: {
        prompt: "Collect warranty information...",
        tools: [recordWarrantyStatus]
      },
      specialist: {
        prompt: `Provide solutions based on warranty: ${request.state.warrantyStatus}`,
        tools: [provideSolution, escalate]
      }
    };

    const config = configs[step];
    request = request.override({  
      systemPrompt: config.prompt,  
      tools: config.tools
    });
    return handler(request);
  }
);

// 4. Create agent with middleware
const agent = createAgent({
  model,
  tools: [recordWarrantyStatus, provideSolution, escalate],
  stateSchema: SupportState,
  middleware: [applyStepConfig],  
  checkpointer: new InMemorySaver()  // Persist state across turns
});

import { StateGraph, START, END, MessagesAnnotation, Annotation, Command } from "@langchain/langgraph";
import { createAgent } from "langchain";
import { tool, ToolRuntime } from "@langchain/core/tools";
import { AIMessage, ToolMessage } from "@langchain/core/messages";
import { z } from "zod";

// 1. Define state with active_agent tracker
const MultiAgentState = Annotation.Root({
  ...MessagesAnnotation.spec,
  activeAgent: Annotation<string | undefined>({
    reducer: (x, y) => y ?? x,
  })
});

// 2. Create handoff tools
const transferToSales = tool(
  async (_, { runtime }: { runtime: ToolRuntime }) => {
    const lastAiMessage = [...runtime.state.messages]  
      .reverse()  
      .find((msg): msg is AIMessage => msg instanceof AIMessage);  
    const transferMessage = new ToolMessage({  
      content: "Transferred to sales agent from support agent",  
      tool_call_id: runtime.toolCallId
    });  
    return new Command({
      goto: "sales_agent",
      update: {
        activeAgent: "sales_agent",
        messages: [lastAiMessage, transferMessage].filter(Boolean),  
      },
      graph: Command.PARENT
    });
  },
  {
    name: "transfer_to_sales",
    description: "Transfer to the sales agent.",
    schema: z.object({})
  }
);

const transferToSupport = tool(
  async (_, { runtime }: { runtime: ToolRuntime }) => {
    const lastAiMessage = [...runtime.state.messages]  
      .reverse()  
      .find((msg): msg is AIMessage => msg instanceof AIMessage);  
    const transferMessage = new ToolMessage({  
      content: "Transferred to support agent from sales agent",  
      tool_call_id: runtime.toolCallId
    });  
    return new Command({
      goto: "support_agent",
      update: {
        activeAgent: "support_agent",
        messages: [lastAiMessage, transferMessage].filter(Boolean),  
      },
      graph: Command.PARENT
    });
  },
  {
    name: "transfer_to_support",
    description: "Transfer to the support agent.",
    schema: z.object({})
  }
);

// 3. Create agents with handoff tools
const salesAgent = createAgent({
  model: "anthropic:claude-sonnet-4-20250514",
  tools: [transferToSupport],
  prompt: "You are a sales agent. Help with sales inquiries. If asked about technical issues or support, transfer to the support agent."
});

const supportAgent = createAgent({
  model: "anthropic:claude-sonnet-4-20250514",
  tools: [transferToSales],
  prompt: "You are a support agent. Help with technical issues. If asked about pricing or purchasing, transfer to the sales agent."
});

// 4. Create agent nodes that invoke the agents
const callSalesAgent = async (state: typeof MultiAgentState.State) => {
  const response = await salesAgent.invoke(state);
  return response;
};

const callSupportAgent = async (state: typeof MultiAgentState.State) => {
  const response = await supportAgent.invoke(state);
  return response;
};

// 5. Create router that checks if we should end or continue
const routeAfterAgent = (
  state: typeof MultiAgentState.State
): "sales_agent" | "support_agent" | "__end__" => {
  const messages = state.messages ?? [];

  // Check the last message - if it's an AIMessage without tool calls, we're done
  if (messages.length > 0) {
    const lastMsg = messages[messages.length - 1];
    if (lastMsg instanceof AIMessage && !lastMsg.tool_calls?.length) {  
      return "__end__";  
    }
  }

  // Otherwise route to the active agent
  const active = state.activeAgent ?? "sales_agent";
  return active as "sales_agent" | "support_agent";
};

const routeInitial = (
  state: typeof MultiAgentState.State
): "sales_agent" | "support_agent" => {
  // Route to the active agent based on state, default to sales agent
  return (state.activeAgent ?? "sales_agent") as "sales_agent" | "support_agent";
};

// 6. Build the graph
const builder = new StateGraph(MultiAgentState)
  .addNode("sales_agent", callSalesAgent)
  .addNode("support_agent", callSupportAgent);

// Start with conditional routing based on initial activeAgent
builder.addConditionalEdges(START, routeInitial, ["sales_agent", "support_agent"]);

// After each agent, check if we should end or route to another agent
builder.addConditionalEdges(
  "sales_agent",
  routeAfterAgent,
  ["sales_agent", "support_agent", END]
);
builder.addConditionalEdges(
  "support_agent",
  routeAfterAgent,
  ["sales_agent", "support_agent", END]
);

const graph = builder.compile();
const result = await graph.invoke({
  messages: [
    {
      role: "user",
      content: "Hi, I'm having trouble with my account login. Can you help?"
    }
  ]
});

for (const msg of result.messages) {
  msg.pretty_print();
}