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

// Load environment variables
dotenv.config();

// Check for required env vars
if (!process.env.OPENAI_API_KEY) {
  console.error("Error: OPENAI_API_KEY environment variable is required");
  process.exit(1);
}

async function runBatchingDemo() {
  console.log("Starting Memory Batching Demo");
  
  // Create Memory instance
  const memory = new Memory({
    embedder: {
      provider: "openai",
      config: {
        apiKey: process.env.OPENAI_API_KEY,
        model: "text-embedding-3-small",
      },
    },
    vectorStore: {
      provider: "memory",
      config: {
        dimension: 1536,
      },
    },
    llm: {
      provider: "openai_structured",
      config: {
        apiKey: process.env.OPENAI_API_KEY,
        model: "gpt-3.5-turbo",
      },
    },
    historyDbPath: ":memory:",
  });
  
  const userId = "batching-demo-user";
  
  // Create a long conversation to test batching
  console.log("\n=== Creating Long Conversation ===");
  
  const longConversation: Message[] = [
    { role: "user", content: "Hi, I'm starting a new software project." },
    { role: "assistant", content: "That's exciting! What kind of software project are you working on?" },
    { role: "user", content: "It's a web application for task management." },
    { role: "assistant", content: "Great choice! What technologies are you planning to use?" },
    { role: "user", content: "I'm thinking React for the frontend and Node.js for the backend." },
    { role: "assistant", content: "Excellent stack! Have you considered which database to use?" },
    { role: "user", content: "I was thinking PostgreSQL for the main database." },
    { role: "assistant", content: "PostgreSQL is a solid choice. What about user authentication?" },
    { role: "user", content: "I'll implement JWT tokens for authentication." },
    { role: "assistant", content: "Good security approach. Any plans for the UI framework?" },
    { role: "user", content: "I'll use Material-UI for the component library." },
    { role: "assistant", content: "Material-UI will give you a clean, professional look. What about state management?" },
    { role: "user", content: "I'm planning to use Redux Toolkit for state management." },
    { role: "assistant", content: "Redux Toolkit is an excellent choice. Have you thought about testing?" },
    { role: "user", content: "Yes, I'll use Jest for unit tests and Cypress for e2e tests." },
  ];
  
  console.log(`Created conversation with ${longConversation.length} messages`);
  
  // Test 1: Default batching (batch size 5, stride 2)
  console.log("\n=== Test 1: Default Batching (size=5, stride=2) ===");
  try {
    const result1 = await memory.add(longConversation, {
      userId,
      infer: true,
      enableBatching: true
    });
    
    console.log(`\n✅ Results: ${result1.results.length} memories created`);
    
    // Show batch information from metadata
    const batchInfo = new Map();
    result1.results.forEach(item => {
      if (item.metadata?.batchIndex) {
        if (!batchInfo.has(item.metadata.batchIndex)) {
          batchInfo.set(item.metadata.batchIndex, {
            batchIndex: item.metadata.batchIndex,
            batchStart: item.metadata.batchStart,
            batchEnd: item.metadata.batchEnd,
            memories: 0
          });
        }
        batchInfo.get(item.metadata.batchIndex).memories++;
      }
    });
    
    console.log("\nBatch breakdown:");
    batchInfo.forEach(batch => {
      console.log(`  Batch ${batch.batchIndex}: messages ${batch.batchStart + 1}-${batch.batchEnd} → ${batch.memories} memories`);
    });
    
  } catch (error) {
    console.error("Error in test 1:", error);
  }
  
  // Reset for next test
  await memory.reset();
  
  // Test 2: Custom batching (batch size 3, stride 1)
  console.log("\n=== Test 2: Custom Batching (size=3, stride=1) ===");
  try {
    const result2 = await memory.add(longConversation, {
      userId,
      infer: true,
      enableBatching: true,
      batchSize: 3,
      batchStride: 1
    });
    
    console.log(`\n✅ Results: ${result2.results.length} memories created`);
    
    // Show batch information
    const batchInfo2 = new Map();
    result2.results.forEach(item => {
      if (item.metadata?.batchIndex) {
        if (!batchInfo2.has(item.metadata.batchIndex)) {
          batchInfo2.set(item.metadata.batchIndex, {
            batchIndex: item.metadata.batchIndex,
            batchStart: item.metadata.batchStart,
            batchEnd: item.metadata.batchEnd,
            memories: 0
          });
        }
        batchInfo2.get(item.metadata.batchIndex).memories++;
      }
    });
    
    console.log("\nBatch breakdown:");
    batchInfo2.forEach(batch => {
      console.log(`  Batch ${batch.batchIndex}: messages ${batch.batchStart + 1}-${batch.batchEnd} → ${batch.memories} memories`);
    });
    
  } catch (error) {
    console.error("Error in test 2:", error);
  }
  
  // Reset for next test
  await memory.reset();
  
  // Test 3: Disabled batching (should process all at once)
  console.log("\n=== Test 3: Disabled Batching ===");
  try {
    const result3 = await memory.add(longConversation, {
      userId,
      infer: true,
      enableBatching: false
    });
    
    console.log(`\n✅ Results: ${result3.results.length} memories created (no batching)`);
    
  } catch (error) {
    console.error("Error in test 3:", error);
  }
  
  // Reset for next test
  await memory.reset();
  
  // Test 4: Small conversation (should not trigger batching)
  console.log("\n=== Test 4: Small Conversation (No Batching Triggered) ===");
  const smallConversation: Message[] = [
    { role: "user", content: "Hello!" },
    { role: "assistant", content: "Hi there! How can I help you?" },
    { role: "user", content: "I need help with my project." },
    { role: "assistant", content: "I'd be happy to help! What's your project about?" },
  ];
  
  try {
    const result4 = await memory.add(smallConversation, {
      userId,
      infer: true,
      enableBatching: true
    });
    
    console.log(`\n✅ Results: ${result4.results.length} memories created (automatic mode, no batching needed)`);
    
  } catch (error) {
    console.error("Error in test 4:", error);
  }
  
  // Test search to see how batched memories work
  console.log("\n=== Testing Search with Batched Memories ===");
  
  // Add the long conversation again with default batching
  await memory.add(longConversation, {
    userId,
    infer: true,
    enableBatching: true
  });
  
  // Search for specific topics
  const searchQueries = [
    "What technologies are being used for the frontend?",
    "Tell me about the database choices",
    "What testing approach is planned?"
  ];
  
  for (const query of searchQueries) {
    console.log(`\n🔍 Searching: "${query}"`);
    const searchResult = await memory.search(query, { userId, limit: 3 });
    
    console.log(`Found ${searchResult.results.length} relevant memories:`);
    searchResult.results.forEach((result, idx) => {
      console.log(`  ${idx + 1}. ${result.memory} (Score: ${result.score?.toFixed(3)})`);
      if (result.metadata?.batchIndex) {
        console.log(`     From batch ${result.metadata.batchIndex} (messages ${result.metadata.batchStart + 1}-${result.metadata.batchEnd})`);
      }
    });
  }
  
  // Test chat with batched memories
  console.log("\n=== Testing Chat with Batched Memories ===");
  const chatQuery = "Summarize the technology stack for the task management project";
  
  try {
    const chatResult = await memory.chat(chatQuery, {
      userId,
      includeContext: true,
      useSearchAgents: true
    });
    
    console.log(`\n💬 Chat Query: "${chatQuery}"`);
    console.log(`Response: ${chatResult.response}`);
    console.log(`Sources used: ${chatResult.metadata.sourcesCount}`);
    
  } catch (error) {
    console.error("Error in chat test:", error);
  }
  
  // Clean up
  console.log("\n=== Cleanup ===");
  await memory.reset();
  console.log("✅ Batching demo completed successfully!");
}

// Run the demo
runBatchingDemo()
  .catch(error => {
    console.error("Error in batching demo:", error);
    process.exit(1);
  });