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 runChatDemo() {
  console.log("Starting Memory Chat 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 = "chat-demo-user";
  
  // Add diverse memories for comprehensive chat testing
  console.log("\n=== Setting up Memory Context ===");
  
  const conversations: Message[][] = [
    [
      { role: "user", content: "Hi, I'm Alex, a software engineer at TechCorp." },
      { role: "assistant", content: "Nice to meet you, Alex! Tell me more about your work at TechCorp." },
      { role: "user", content: "I work on machine learning projects, specifically computer vision." },
    ],
    [
      { role: "user", content: "I'm planning a trip to Japan next month." },
      { role: "assistant", content: "That sounds exciting! What places are you planning to visit?" },
      { role: "user", content: "Tokyo and Kyoto mainly. I'm interested in traditional culture and modern tech." },
    ],
    [
      { role: "user", content: "My colleague Sarah recommended a great book about AI ethics." },
      { role: "user", content: "It's called 'The Ethical Algorithm' by Kearns and Roth." },
    ],
    [
      { role: "user", content: "Had a productive meeting with the product team today." },
      { role: "user", content: "We discussed the new computer vision feature for our mobile app." },
      { role: "user", content: "Sarah thinks we can have a prototype ready by next quarter." },
    ],
  ];
  
  // Add the conversations to memory
  for (let i = 0; i < conversations.length; i++) {
    const conv = conversations[i];
    console.log(`Adding conversation ${i + 1}...`);
    await memory.add(conv, { userId, infer: true });
  }
  
  // Additional individual memories
  const additionalMemories = [
    "Alex's favorite programming language is Python.",
    "TechCorp is located in San Francisco.",
    "Sarah is the project manager for the AI initiative.",
    "The Japan trip is scheduled for March 2024.",
  ];
  
  for (const mem of additionalMemories) {
    await memory.add(mem, { userId, infer: true });
  }
  
  console.log("✅ Memory context established\n");
  
  // Test various chat scenarios
  const chatQueries = [
    {
      query: "Tell me about Alex",
      options: { userId, includeContext: true, useSearchAgents: false },
      description: "Basic chat without agents"
    },
    {
      query: "What do you know about Alex's work?",
      options: { userId, includeContext: true, useSearchAgents: true },
      description: "Enhanced chat with agents"
    },
    {
      query: "Tell me about Alex's upcoming travel plans",
      options: { userId, includeContext: false, useSearchAgents: false },
      description: "Chat without context messages"
    },
    {
      query: "What are Alex and Sarah working on together?",
      options: { userId, includeContext: true, useSearchAgents: true, maxAgentSteps: 3 },
      description: "Complex query with enhanced agents"
    },
    {
      query: "What programming languages does Alex use?",
      options: { userId, includeContext: true, useSearchAgents: false },
      description: "Specific fact retrieval"
    },
    {
      query: "Tell me about the team dynamics at TechCorp",
      options: { userId, includeContext: true, useSearchAgents: true },
      description: "Inferential reasoning chat"
    },
  ];
  
  for (let i = 0; i < chatQueries.length; i++) {
    const { query, options, description } = chatQueries[i];
    
    console.log(`\n=== Chat Test ${i + 1}: ${description} ===`);
    console.log(`Query: "${query}"`);
    console.log(`Options: includeContext=${options.includeContext}, useSearchAgents=${options.useSearchAgents}`);
    
    try {
      const startTime = Date.now();
      const chatResult = await memory.chat(query, options);
      const responseTime = Date.now() - startTime;
      
      console.log(`\n📝 Response (${responseTime}ms):`);
      console.log(chatResult.response);
      
      console.log(`\n📊 Metadata:`);
      console.log(`- Sources used: ${chatResult.metadata.sourcesCount}`);
      console.log(`- Response tokens: ~${chatResult.metadata.responseTokens}`);
      console.log(`- Processing time: ${chatResult.metadata.processingTimeMs}ms`);
      
      if (chatResult.sources.length > 0) {
        console.log(`\n📚 Sources:`);
        chatResult.sources.slice(0, 3).forEach((source, idx) => {
          console.log(`${idx + 1}. ${source.memory} (Score: ${source.score?.toFixed(3)})`);
          if (source.entityIds && source.entityIds.length > 0) {
            console.log(`   Entities: ${source.entityIds.join(', ')}`);
          }
        });
      }
      
      if (chatResult.agentAnalysis) {
        console.log(`\n🤖 Agent Analysis:`);
        if (chatResult.agentAnalysis.decision) {
          console.log(`- Additional steps needed: ${chatResult.agentAnalysis.decision.needsAdditionalSteps}`);
          if (chatResult.agentAnalysis.decision.needsAdditionalSteps) {
            console.log(`- Reasoning: ${chatResult.agentAnalysis.decision.reasoning}`);
            console.log(`- Steps executed: ${chatResult.agentAnalysis.stepResults?.length || 0}`);
            console.log(`- Additional results: ${chatResult.agentAnalysis.additionalResults?.length || 0}`);
          }
        } else if (chatResult.agentAnalysis.isMultihop && chatResult.agentAnalysis.multihopResults) {
          console.log(`- Multihop analysis with ${chatResult.agentAnalysis.multihopResults.length} hops`);
          chatResult.agentAnalysis.multihopResults.forEach(hop => {
            console.log(`  Hop ${hop.hopNumber}: ${hop.decision.reasoning}`);
          });
        }
      }
      
    } catch (error) {
      console.error(`❌ Error in chat test ${i + 1}:`, error);
    }
  }
  
  // Test edge cases
  console.log(`\n=== Edge Case Tests ===`);
  
  // Query with no relevant memories
  console.log(`\n🔍 Testing query with no relevant memories:`);
  const noResultsQuery = "What do you know about quantum computing?";
  try {
    const result = await memory.chat(noResultsQuery, { userId });
    console.log(`Query: "${noResultsQuery}"`);
    console.log(`Response: ${result.response}`);
    console.log(`Sources found: ${result.metadata.sourcesCount}`);
  } catch (error) {
    console.error("Error in no-results test:", error);
  }
  
  // High limit test
  console.log(`\n🔍 Testing with high limit:`);
  try {
    const result = await memory.chat("Tell me everything about Alex", { 
      userId, 
      limit: 50, 
      includeContext: true, 
      useSearchAgents: true 
    });
    console.log(`Query: "Tell me everything about Alex"`);
    console.log(`Sources found: ${result.metadata.sourcesCount}`);
    console.log(`Response length: ${result.response.length} characters`);
  } catch (error) {
    console.error("Error in high-limit test:", error);
  }
  
  // Clean up
  console.log("\n=== Cleanup ===");
  await memory.reset();
  console.log("✅ Chat demo completed successfully!");
}

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