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

// 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 runMemoryExample() {
  console.log("Starting Memory Example with Entity Extraction");
  
  // Create Memory instance with OpenAI embeddings and in-memory vector store
  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, // Dimensions for OpenAI's text-embedding-3-small
      },
    },
    llm: {
      provider: "openai_structured", // Using the structured OpenAI LLM
      config: {
        apiKey: process.env.OPENAI_API_KEY,
        model: "gpt-3.5-turbo",
      },
    },
    historyDbPath: ":memory:", // Use in-memory SQLite database
  });
  
  // User ID for this example
  const userId = "user-123";
  
  // Add some memories with entity extraction
  console.log("\n=== Adding Memories with Entity Extraction ===");
  
  const messages: Message[] = [
    { role: "user", content: "Hi, my name is Saket. I am a software engineer." },
    { role: "user", content: "I live in India and work at Google." },
    { role: "user", content: "I am going to San Francisco next month for vacation." },
    { role: "user", content: "John and I had a meeting yesterday about the new AI project." },
  ];
  
  const addResult = await memory.add(messages, {
    userId,
    infer: true, // Use LLM to extract facts and entities from messages
  });
  
  console.log(`Added ${addResult.results.length} memories with entity extraction:`);
  addResult.results.forEach(item => {
    console.log(`- ${item.memory} (ID: ${item.id.slice(0, 8)}...)`);
    if (item.entityIds && item.entityIds.length > 0) {
      console.log(`  Associated Entities: ${item.entityIds.join(', ')}`);
    }
    console.log(`  Event: ${item.metadata?.event}`);
  });

  // View all extracted entities
  console.log("\n=== All Extracted Entities ===");
  const entities = await memory.getAllEntities({ userId });
  console.log(`Found ${entities.length} entities:`);
  entities.forEach(entity => {
    console.log(`- ${entity.entity_id}: ${entity.label} (${entity.type})`);
  });

  // Add more memories to test entity association
  console.log("\n=== Adding More Memories ===");
  const additionalMessages = [
    "Saket's favorite programming language is TypeScript.",
    "Google has an office in San Francisco.",
    "John is the project manager for the AI initiative.",
    "I met Sarah from the design team at Google."
  ];
  
  for (const msg of additionalMessages) {
    const result = await memory.add(msg, { 
      userId,
      infer: true // Enable entity extraction
    });
    
    console.log(`Added: "${msg}"`);
    if (result.results[0]?.entityIds && result.results[0].entityIds.length > 0) {
      console.log(`  Entities: ${result.results[0].entityIds.join(', ')}`);
    }
  }

  // Search for memories with entity information
  console.log("\n=== Searching Memories with Entity Information ===");
  const searchQuery = "Tell me about Saket";
  console.log(`Searching for: "${searchQuery}"`);
  
  const searchResult = await memory.search(searchQuery, { 
    userId,
    includeContext: true,
    contextWindowSize: 3 // Get 3 messages before and after each memory
  });
  
  console.log(`Found ${searchResult.results.length} memories:`);
  searchResult.results.forEach(item => {
    console.log(`- ${item.memory} (Score: ${item.score?.toFixed(3)})`);
    if (item.entityIds && item.entityIds.length > 0) {
      console.log(`  Associated Entities: ${item.entityIds.join(', ')}`);
    }
    
    // Display original messages if available
    if (item.messages && item.messages.length > 0) {
      console.log("  Original messages:");
      item.messages.forEach((msg, idx) => {
        console.log(`    ${idx + 1}. [${msg.role}]: ${typeof msg.content === 'string' ? msg.content : '[multimodal content]'}`);
      });
    }
    
    // Display context if available
    if (item.context && item.context.length > 0) {
      console.log("  Context messages:");
      item.context.forEach(msg => {
        const typeSymbol = msg.type === 'before' ? '↑' : msg.type === 'after' ? '↓' : '→';
        console.log(`  ${typeSymbol} ${(msg.content || '[No content]').substring(0, 50)}${(msg.content || '').length > 50 ? '...' : ''}`);
      });
    }
  });

  // Search for entity-specific information
  console.log("\n=== Entity-Specific Search ===");
  const entityQueries = [
    "What do you know about Google?",
    "Tell me about John",
    "What about San Francisco?"
  ];

  for (const query of entityQueries) {
    console.log(`\nSearching: "${query}"`);
    const result = await memory.search(query, { userId, limit: 3 });
    result.results.forEach(item => {
      console.log(`  - ${item.memory} (Score: ${item.score?.toFixed(3)})`);
      if (item.entityIds && item.entityIds.length > 0) {
        console.log(`    Entities: ${item.entityIds.join(', ')}`);
      }
    });
  }

  // Test entity management operations
  console.log("\n=== Entity Management Operations ===");
  
  // Add a custom entity
  await memory.addEntity("typescript", "TypeScript", "Concept", userId);
  console.log("Added custom entity: TypeScript (Concept)");
  
  // Get specific entity
  const saketEntity = await memory.getEntity("saket");
  if (saketEntity) {
    console.log(`Retrieved entity: ${saketEntity.entity_id} - ${saketEntity.label} (${saketEntity.type})`);
  }
  
  // Update an entity
  await memory.updateEntity("john", "John Smith", "Person", userId);
  console.log("Updated John entity with full name");

  // Test memory-entity associations
  const memoryId = addResult.results[0].id;
  console.log(`\nMemory-Entity Associations for memory: ${memoryId.slice(0, 8)}...`);
  const memoryEntities = await memory.getMemoryEntities(memoryId);
  console.log(`Associated entities: ${memoryEntities.join(', ')}`);

  // Add a new memory with entity extraction
  console.log("\n=== Adding Complex Memory ===");
  const complexMemory = "During the meeting in Google's San Francisco office, Saket and John discussed implementing the new TypeScript features for the AI project with Sarah's design input.";
  const complexResult = await memory.add(complexMemory, { 
    userId,
    infer: true 
  });
  
  console.log(`Added complex memory: "${complexMemory}"`);
  if (complexResult.results[0]?.entityIds) {
    console.log(`Extracted entities: ${complexResult.results[0].entityIds.join(', ')}`);
  }

  // Get all memories with entity information
  console.log("\n=== All Memories with Entities ===");
  const allMemories = await memory.getAll({ userId });
  console.log(`Total memories: ${allMemories.results.length}`);
  allMemories.results.forEach((item, index) => {
    console.log(`${index + 1}. ${item.memory}`);
    if (item.entityIds && item.entityIds.length > 0) {
      console.log(`   Entities: ${item.entityIds.join(', ')}`);
    }
    if (item.messages && item.messages.length > 0) {
      console.log(`   Original messages: ${item.messages.length} message(s)`);
      item.messages.forEach((msg, idx) => {
        const content = typeof msg.content === 'string' ? msg.content.substring(0, 60) : '[multimodal]';
        console.log(`     ${idx + 1}. [${msg.role}]: ${content}${content.length > 60 ? '...' : ''}`);
      });
    }
  });
  
  // Update all entities to see the current state
  console.log("\n=== Final Entity State ===");
  const finalEntities = await memory.getAllEntities({ userId });
  console.log(`Total entities: ${finalEntities.length}`);
  finalEntities.forEach(entity => {
    console.log(`- ${entity.entity_id}: ${entity.label} (${entity.type})`);
  });

  // View history of a memory with entities
  console.log("\n=== Memory History ===");
  const history = await memory.history(memoryId, { userId });
  console.log(`History events for memory: ${history.length}`);
  history.forEach(event => {
    console.log(`- ${event.action}: ${event.new_value || "[deleted]"} (${event.happened_at})`);
  });
  
  // Get all messages
  console.log("\n=== Recent Messages ===");
  const allMessages = await memory.getAllMessages({ userId, limit: 5 });
  console.log(`Recent messages: ${allMessages.length}`);
  allMessages.forEach(message => {
    console.log(`- ${message.action}: ${message.new_value || "[deleted]"} (${message.happened_at})`);
  });

  // Test deleting an entity
  console.log("\n=== Entity Deletion Test ===");
  console.log("Deleting 'typescript' entity...");
  await memory.deleteEntity("typescript");
  
  const entitiesAfterDeletion = await memory.getAllEntities({ userId });
  console.log(`Entities after deletion: ${entitiesAfterDeletion.length}`);
  entitiesAfterDeletion.forEach(entity => {
    console.log(`- ${entity.entity_id}: ${entity.label} (${entity.type})`);
  });

  // Clean up - reset memory for this example
  console.log("\n=== Resetting Memory Store ===");
  await memory.reset();
  console.log("Memory store and entities reset successfully");
  
  console.log("\nMemory Example with Entity Extraction Completed");
}

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