import { Memory } from "../../src/memory";
import { Message } from "../../src/types";

async function runSearchAgentsExample() {
  console.log("Starting Memory Search Agents Example");

  const memory = new Memory({
    llm: {
      provider: "openai",
      config: {
        apiKey: process.env.OPENAI_API_KEY,
        model: "gpt-4o-mini",
      },
    },
    embedder: {
      provider: "openai",
      config: {
        apiKey: process.env.OPENAI_API_KEY,
        model: "text-embedding-3-small",
      },
    },
  });

  const userId = "search-agent-user";

  // First, let's add some sequential memories with timestamps
  console.log("\n=== Adding Sequential Memories ===");
  
  const sequentialEvents = [
    { content: "Started working on the AI project with John", delay: 0 },
    { content: "Had initial brainstorming session about machine learning features", delay: 100 },
    { content: "John suggested using TypeScript for better type safety", delay: 200 },
    { content: "Decided to integrate with the vector database", delay: 300 },
    { content: "Sarah joined the project as a designer", delay: 400 },
    { content: "Created wireframes for the user interface", delay: 500 },
    { content: "Implemented the first prototype", delay: 600 },
    { content: "Got feedback from the design team", delay: 700 },
    { content: "Made improvements based on Sarah's suggestions", delay: 800 },
    { content: "Project completed successfully", delay: 900 },
  ];

  for (let i = 0; i < sequentialEvents.length; i++) {
    const event = sequentialEvents[i];
    
    // Add artificial delay to create temporal separation
    if (event.delay > 0) {
      await new Promise(resolve => setTimeout(resolve, event.delay));
    }
    
    await memory.add(event.content, {
      userId,
      infer: true
    });
    
    console.log(`${i + 1}. Added: "${event.content}"`);
  }

  // Test 1: Basic search without agents
  console.log("\n=== Test 1: Basic Search (No Agents) ===");
  const basicSearch = await memory.search("What happened with John and the project?", {
    userId,
    useSearchAgents: false,
    limit: 3
  });

  console.log(`Found ${basicSearch.results.length} basic results:`);
  basicSearch.results.forEach((result, i) => {
    console.log(`${i + 1}. ${result.memory} (Score: ${result.score?.toFixed(3)})`);
    if (result.entityIds?.length) {
      console.log(`   Entities: ${result.entityIds.join(', ')}`);
    }
  });

  // Test 2: Enhanced search with agents
  console.log("\n=== Test 2: Enhanced Search (With Agents) ===");
  const enhancedSearch = await memory.search("What happened with John and the project?", {
    userId,
    useSearchAgents: true,
    maxAgentSteps: 2,
    limit: 3
  });

  console.log(`Found ${enhancedSearch.results.length} initial results:`);
  enhancedSearch.results.forEach((result, i) => {
    console.log(`${i + 1}. ${result.memory} (Score: ${result.score?.toFixed(3)})`);
    if (result.entityIds?.length) {
      console.log(`   Entities: ${result.entityIds.join(', ')}`);
    }
  });

  if (enhancedSearch.agentAnalysis) {
    console.log("\n--- Agent Analysis ---");
    if (enhancedSearch.agentAnalysis.decision) {
      console.log(`Needs additional steps: ${enhancedSearch.agentAnalysis.decision.needsAdditionalSteps}`);
      console.log(`Reasoning: ${enhancedSearch.agentAnalysis.decision.reasoning}`);
      
      if (enhancedSearch.agentAnalysis.decision.nextSteps.length > 0) {
        console.log("\nPlanned steps:");
        enhancedSearch.agentAnalysis.decision.nextSteps.forEach((step, i) => {
          console.log(`${i + 1}. ${step.action} (Priority: ${step.priority})`);
          if (Object.keys(step.parameters).length > 0) {
            console.log(`   Parameters: ${JSON.stringify(step.parameters)}`);
          }
        });
      }
    }

    if (enhancedSearch.agentAnalysis.stepResults && enhancedSearch.agentAnalysis.stepResults.length > 0) {
      console.log("\nExecuted steps:");
      enhancedSearch.agentAnalysis.stepResults.forEach((result, i) => {
        console.log(`${i + 1}. ${result.action}: ${result.success ? 'SUCCESS' : 'FAILED'}`);
        if (result.results && Array.isArray(result.results)) {
          console.log(`   Found ${result.results.length} additional items`);
        }
      });
    }

    if (enhancedSearch.agentAnalysis.additionalResults && enhancedSearch.agentAnalysis.additionalResults.length > 0) {
      console.log(`\nAdditional results found: ${enhancedSearch.agentAnalysis.additionalResults.length}`);
      enhancedSearch.agentAnalysis.additionalResults.slice(0, 3).forEach((result, i) => {
        console.log(`${i + 1}. [${result.type}] ${result.content}`);
      });
    }
  }

  // Test 3: Temporal query
  console.log("\n=== Test 3: Temporal Query ===");
  const temporalSearch = await memory.search("What happened after John suggested TypeScript?", {
    userId,
    useSearchAgents: true,
    maxAgentSteps: 3
  });

  console.log(`Temporal search results: ${temporalSearch.results.length}`);
  if (temporalSearch.agentAnalysis) {
    console.log(`Agent reasoning: ${temporalSearch.agentAnalysis.decision?.reasoning || 'No decision reasoning'}`);
    console.log(`Additional results: ${temporalSearch.agentAnalysis.additionalResults?.length || 0}`);
  }

  // Test 4: Entity-focused query
  console.log("\n=== Test 4: Entity-Focused Query ===");
  const entitySearch = await memory.search("Tell me everything about Sarah's involvement", {
    userId,
    useSearchAgents: true,
    maxAgentSteps: 2
  });

  console.log(`Entity search results: ${entitySearch.results.length}`);
  if (entitySearch.agentAnalysis) {
    console.log(`Agent reasoning: ${entitySearch.agentAnalysis.decision?.reasoning || 'No decision reasoning'}`);
    if (entitySearch.agentAnalysis.stepResults && entitySearch.agentAnalysis.stepResults.length > 0) {
      console.log("Agent actions taken:");
      entitySearch.agentAnalysis.stepResults.forEach(step => {
        console.log(`- ${step.action}: ${step.success ? 'SUCCESS' : 'FAILED'}`);
      });
    }
  }

  // Test 5: Time-based query
  console.log("\n=== Test 5: Time-Based Query ===");
  const timeSearch = await memory.search("What happened today in the project?", {
    userId,
    useSearchAgents: true,
    maxAgentSteps: 2
  });

  console.log(`Time-based search results: ${timeSearch.results.length}`);
  if (timeSearch.agentAnalysis) {
    console.log(`Agent reasoning: ${timeSearch.agentAnalysis.decision?.reasoning || 'No decision reasoning'}`);
    console.log(`Steps planned: ${timeSearch.agentAnalysis.decision?.nextSteps?.length || 0}`);
  }

  // Clean up
  console.log("\n=== Cleanup ===");
  await memory.reset();
  console.log("Search Agents Example completed successfully!");
}

// Run the example
runSearchAgentsExample()
  .catch(error => {
    console.error("Error in search agents example:", error);
    process.exit(1);
  });