import {
  addAgent,
  addCompletion,
  addSwarm,
  addTool,
  changeAgent,
  execute,
  session,
  Adapter
} from "agent-swarm-kit";

const NAVIGATE_TOOL = addTool({
  toolName: "navigate-tool",
  call: async (clientId, agentName, { to }) => {
    await changeAgent(to, clientId);
    await execute("Navigation complete. Notify the user", clientId);
  },
  validate: async () => true,
  type: "function",
  function: {
    name: "navigate-tool",
    description: "The tool for navigation",
    parameters: {
      type: "object",
      properties: {
        to: {
          type: "string",
          description: "The target agent for navigation",
        },
      },
      required: ["to"],
    },
  },
});

const ollama = new Ollama({ host: CC_OLLAMA_HOST });

const MOCK_COMPLETION = addCompletion({
  completionName: "navigate-completion",
  /**
   * Use whatever you want: NVIDIA NIM, OpenAI, GPT4All, Ollama or LM Studio
   * Even mock it for unit test of tool integration like it done in `test` folder
   * 
   * @see https://github.com/tripolskypetr/agent-swarm-kit/tree/master/test
   */
  getCompletion: Adapter.fromOllama(ollama, "nemotron-mini:4b"), // "tripolskypetr/gemma3-tools:4b"
});

const TRIAGE_AGENT = addAgent({
  agentName: "triage-agent",
  completion: MOCK_COMPLETION,
  prompt: "You are to triage a users request, and call a tool to transfer to the right agent. There are two agents available: `sales-agent` and `refund-agent`",
  tools: [NAVIGATE_TOOL],
});

const SALES_AGENT = addAgent({
  agentName: "sales-agent",
  completion: MOCK_COMPLETION,
  prompt: "You are a sales agent that handles all actions related to placing the order to purchase an item.",
  tools: [NAVIGATE_TOOL],
});

const REDUND_AGENT = addAgent({
  agentName: "refund-agent",
  completion: MOCK_COMPLETION,
  prompt: "You are a refund agent that handles all actions related to refunds after a return has been processed.",
  tools: [NAVIGATE_TOOL],
});

const TEST_SWARM = addSwarm({
  agentList: [TRIAGE_AGENT, SALES_AGENT, REDUND_AGENT],
  defaultAgent: TRIAGE_AGENT,
  swarmName: "navigation-swarm",
});

...

app.get("/api/v1/session/:clientId", upgradeWebSocket((ctx) => {
  const clientId = ctx.req.param("clientId");

  const { complete, dispose } = session(clientId, TEST_SWARM)

  return {
    onMessage: async (event, ws) => {
      const message = event.data.toString();
      ws.send(await complete(message));
    },
    onClose: async () => {
      await dispose();
    },
  }
}));

export enum ToolName {
  TestTool = "test-tool",
}

export enum AgentName {
  TestAgent = "test-agent",
}

export enum CompletionName {
  TestCompletion = "test-completion"
}

export enum SwarmName {
  TestSwarm = "test-swarm"
}

...

addTool({
    toolName: ToolName.TestTool
    ...
})

addAgent({
  agentName: AgentName.TestAgent,
  completion: CompletionName.TestCompletion,
  prompt: "...",
  tools: [ToolName.TestTool],
});

addSwarm({
  agentList: [AgentName.TestAgent],
  defaultAgent: AgentName.TestAgent,
  swarmName: SwarmName.TestSwarm,
});

const { complete, dispose } = session(clientId, SwarmName.TestSwarm)

complete("I need a refund!").then(console.log);

import { addAgent, addSwarm, addStorage, addEmbedding, session, Storage } from "agent-swarm-kit";
import { Ollama } from "ollama";

const ollama = new Ollama();

// Define an embedding with similarity calculation
const NOMIC_EMBEDDING = addEmbedding({
  embeddingName: "nomic_embedding",
  calculateSimilarity: (a, b) => {
    return tidy(() => {
      const tensorA = tensor1d(a);
      const tensorB = tensor1d(b);
      const dotProduct = sum(mul(tensorA, tensorB));
      const normA = norm(tensorA);
      const normB = norm(tensorB);
      const cosineData = div(dotProduct, mul(normA, normB)).dataSync();
      const cosineSimilarity = cosineData[0];
      return cosineSimilarity;
    });
  },
  createEmbedding: async (text) => {
    const { embedding } = await ollama.embeddings({
      model: "nomic-embed-text",
      prompt: text,
    });
    return embedding;
  },
});

// Set up storage with sample data
const TEST_STORAGE = addStorage({
  storageName: "test_storage",
  embedding: NOMIC_EMBEDDING,
  createIndex: ({ description }) => description,
  getData: () => JSON.parse(readFileSync("./data.json", "utf8")).slice(0, 100),
});

// Create an agent with storage
const TEST_AGENT = addAgent({
  agentName: "test_agent",
  completion: TEST_COMPLETION
  prompt: "...",
  storages: [TEST_STORAGE],
});

// Build the swarm
const TEST_SWARM = addSwarm({
  swarmName: "test_swarm",
  agentList: [TEST_AGENT],
  defaultAgent: TEST_AGENT,
});

// Talk with a client
const { complete } = session("client-123", TEST_SWARM);
complete("I need a refund!").then(console.log);

...

export interface PhoneModel {
    id: number;
    title: string;
    description: string;
    diagonal: number;
}

// Use vector search in a tool call
Storage.take<PhoneModel>({
  search: "8 inch phone",
  agentName: AgentName.TestAgent,
  clientId,
  storageName: StorageName.PhoneStorage,
  total: 1,
  score: 0.68,
}).then((phones) => console.log(phones));