import { openai } from '@ai-sdk/openai';
import { streamText, tool } from 'ai';
import { z } from 'zod';
 
// Define tools for the AI to use
const tools = {
  searchKnowledgeBase: tool({
    description: 'Search the knowledge base for relevant information',
    parameters: z.object({
      query: z.string().describe('The search query'),
      maxResults: z.number().optional().describe('Maximum number of results'),
    }),
    execute: async ({ query, maxResults = 5 }) => {
      const results = await vectorStore.search(query, { limit: maxResults });
      return results.map(r => ({
        content: r.content,
        source: r.metadata.source,
        relevance: r.score,
      }));
    },
  }),
 
  createDocument: tool({
    description: 'Create a new document in the system',
    parameters: z.object({
      title: z.string(),
      content: z.string(),
      category: z.enum(['note', 'task', 'reference']),
    }),
    execute: async ({ title, content, category }) => {
      const doc = await documentService.create({ title, content, category });
      return { id: doc.id, message: `Document "${title}" created successfully` };
    },
  }),
 
  queryDatabase: tool({
    description: 'Query the application database',
    parameters: z.object({
      query: z.string().describe('Natural language query to translate to SQL'),
    }),
    execute: async ({ query }) => {
      const sql = await generateSQL(query);
      const results = await db.execute(sql);
      return { rowCount: results.length, data: results.slice(0, 10) };
    },
  }),
};
 
// Stream text with tools
export async function POST(req: Request) {
  const { messages, conversationId } = await req.json();
 
  // Load conversation history
  const history = await conversationStore.getHistory(conversationId);
 
  // Retrieve relevant context for the latest message
  const lastMessage = messages[messages.length - 1];
  const context = await retrieveContext(lastMessage.content);
 
  const result = streamText({
    model: openai('gpt-4o'),
    system: `You are a helpful assistant with access to a knowledge base and tools.
    
Context from knowledge base:
${context.map(c => `- ${c.content} (source: ${c.source})`).join('\n')}
 
Always ground your responses in the provided context when available.`,
    messages: [...history, ...messages],
    tools,
    maxSteps: 5,  // Allow multi-step tool calling
    onStepFinish: async ({ toolCalls, toolResults }) => {
      // Log tool usage for analytics
      if (toolCalls) {
        await analytics.trackToolUsage(conversationId, toolCalls);
      }
    },
    onFinish: async ({ text, usage }) => {
      // Save conversation history
      await conversationStore.addMessages(conversationId, [
        { role: 'user', content: lastMessage.content },
        { role: 'assistant', content: text, toolCalls: usage },
      ]);
      // Track token usage for billing
      await billing.trackUsage(usage);
    },
  });
 
  return result.toDataStreamResponse();
}

import { openai } from '@ai-sdk/openai';
import { embed } from 'ai';
 
interface Document {
  id: string;
  content: string;
  metadata: {
    source: string;
    title: string;
    chunkIndex: number;
  };
}
 
class VectorStore {
  private documents: Document[] = [];
  private embeddings: Map<string, number[]> = new Map();
 
  async addDocument(doc: Document): Promise<void> {
    const embedding = await this.generateEmbedding(doc.content);
    this.documents.push(doc);
    this.embeddings.set(doc.id, embedding);
  }
 
  async search(query: string, options: { limit?: number; threshold?: number } = {}): Promise<Array<Document & { score: number }>> {
    const { limit = 5, threshold = 0.7 } = options;
    const queryEmbedding = await this.generateEmbedding(query);
 
    const results = this.documents
      .map(doc => ({
        ...doc,
        score: this.cosineSimilarity(queryEmbedding, this.embeddings.get(doc.id)!),
      }))
      .filter(r => r.score >= threshold)
      .sort((a, b) => b.score - a.score)
      .slice(0, limit);
 
    return results;
  }
 
  private async generateEmbedding(text: string): Promise<number[]> {
    const { embedding } = await embed({
      model: openai.embedding('text-embedding-3-small'),
      value: text,
    });
    return embedding;
  }
 
  private cosineSimilarity(a: number[], b: number[]): number {
    let dotProduct = 0;
    let normA = 0;
    let normB = 0;
    for (let i = 0; i < a.length; i++) {
      dotProduct += a[i] * b[i];
      normA += a[i] * a[i];
      normB += b[i] * b[i];
    }
    return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
  }
}
 
// Document chunking strategy
function chunkDocument(content: string, maxChunkSize = 1000, overlap = 200): string[] {
  const chunks: string[] = [];
  const sentences = content.split(/(?<=[.!?])\s+/);
  let currentChunk = '';
 
  for (const sentence of sentences) {
    if (currentChunk.length + sentence.length > maxChunkSize && currentChunk.length > 0) {
      chunks.push(currentChunk.trim());
      // Keep overlap from end of previous chunk
      const words = currentChunk.split(' ');
      currentChunk = words.slice(-Math.floor(overlap / 5)).join(' ') + ' ' + sentence;
    } else {
      currentChunk += (currentChunk ? ' ' : '') + sentence;
    }
  }
 
  if (currentChunk.trim()) {
    chunks.push(currentChunk.trim());
  }
 
  return chunks;
}

class ConversationMemory {
  private conversations: Map<string, Message[]> = new Map();
  private summaries: Map<string, string> = new Map();
  private maxMessages = 50;
 
  async addMessage(conversationId: string, message: Message): Promise<void> {
    const history = this.conversations.get(conversationId) || [];
    history.push(message);
 
    // Summarize old messages when history gets too long
    if (history.length > this.maxMessages) {
      await this.summarizeOldMessages(conversationId, history);
    }
 
    this.conversations.set(conversationId, history);
  }
 
  async getContext(conversationId: string, maxTokens: number = 4000): Promise<Message[]> {
    const history = this.conversations.get(conversationId) || [];
    const summary = this.summaries.get(conversationId);
 
    const context: Message[] = [];
 
    // Add summary if available
    if (summary) {
      context.push({
        role: 'system',
        content: `Previous conversation summary: ${summary}`,
      });
    }
 
    // Add recent messages that fit within token limit
    let tokenCount = summary ? this.estimateTokens(summary) : 0;
    for (let i = history.length - 1; i >= 0; i--) {
      const messageTokens = this.estimateTokens(history[i].content);
      if (tokenCount + messageTokens > maxTokens) break;
      tokenCount += messageTokens;
      context.unshift(history[i]);
    }
 
    return context;
  }
 
  private async summarizeOldMessages(conversationId: string, messages: Message[]): Promise<void> {
    const oldMessages = messages.slice(0, messages.length - this.maxMessages);
    const existingSummary = this.summaries.get(conversationId) || '';
 
    const result = await generateText({
      model: openai('gpt-4o-mini'),
      prompt: `Summarize this conversation concisely, preserving key facts and decisions:
${existingSummary ? `Previous summary: ${existingSummary}\n\n` : ''}
${oldMessages.map(m => `${m.role}: ${m.content}`).join('\n')}`,
    });
 
    this.summaries.set(conversationId, result.text);
 
    // Remove summarized messages from history
    const remaining = messages.slice(messages.length - this.maxMessages);
    this.conversations.set(conversationId, remaining);
  }
 
  private estimateTokens(text: string): number {
    return Math.ceil(text.length / 4);
  }
}

interface EvaluationResult {
  score: number;
  reasoning: string;
  criteria: string;
}
 
class AIEvaluator {
  async evaluateResponse(
    query: string,
    response: string,
    context?: string[],
    groundTruth?: string
  ): Promise<{
    relevance: EvaluationResult;
    accuracy: EvaluationResult;
    completeness: EvaluationResult;
    overall: number;
  }> {
    const evalPrompt = `Evaluate this AI response:
 
Query: ${query}
Response: ${response}
${context ? `Context: ${context.join('\n')}` : ''}
${groundTruth ? `Ground Truth: ${groundTruth}` : ''}
 
Rate each criterion from 0-10 and provide reasoning:
1. Relevance: Does the response address the query?
2. Accuracy: Is the information correct?
3. Completeness: Does it cover all aspects of the query?
 
Respond in JSON format.`;
 
    const result = await generateText({
      model: openai('gpt-4o'),
      prompt: evalPrompt,
    });
 
    const evaluation = JSON.parse(result.text);
    const overall = (evaluation.relevance.score + evaluation.accuracy.score + evaluation.completeness.score) / 3;
 
    return { ...evaluation, overall };
  }
}

// Multi-tier caching for AI responses
class AICache {
  private l1Cache: Map<string, { response: string; timestamp: number }> = new Map();
  private l2Cache: Redis;
 
  constructor(redis: Redis) {
    this.l2Cache = redis;
  }
 
  async get(key: string): Promise<string | null> {
    // Check L1 (in-memory) cache
    const l1 = this.l1Cache.get(key);
    if (l1 && Date.now() - l1.timestamp < 60000) { // 1 minute
      return l1.response;
    }
 
    // Check L2 (Redis) cache
    const l2 = await this.l2Cache.get(key);
    if (l2) {
      this.l1Cache.set(key, { response: l2, timestamp: Date.now() });
      return l2;
    }
 
    return null;
  }
 
  async set(key: string, response: string, ttlSeconds: number = 3600): Promise<void> {
    this.l1Cache.set(key, { response, timestamp: Date.now() });
    await this.l2Cache.setex(key, ttlSeconds, response);
  }
 
  generateCacheKey(messages: Message[], model: string): string {
    const content = messages.map(m => `${m.role}:${m.content}`).join('|');
    return `ai:${model}:${hashContent(content)}`;
  }
}

class MultiAgentSystem {
  private agents: Map<string, Agent> = new Map();
 
  registerAgent(name: string, agent: Agent): void {
    this.agents.set(name, agent);
  }
 
  async process(query: string, context: ConversationContext): Promise<string> {
    // Router determines which agent handles the query
    const routerResult = await this.routeToAgent(query);
 
    const agent = this.agents.get(routerResult.agentName);
    if (!agent) {
      return 'I cannot handle this type of request.';
    }
 
    // Execute the selected agent
    const result = await agent.execute(query, context);
 
    // Log for evaluation
    await this.logRouting(query, routerResult, result);
 
    return result;
  }
 
  private async routeToAgent(query: string): Promise<{ agentName: string; confidence: number }> {
    const result = await generateText({
      model: openai('gpt-4o-mini'),
      prompt: `Classify this query into one of these categories: ${Array.from(this.agents.keys()).join(', ')}
 
Query: ${query}
 
Respond with JSON: {"agentName": "...", "confidence": 0.0-1.0}`,
    });
 
    return JSON.parse(result.text);
  }
}

describe('AI Application', () => {
  it('should stream responses token by token', async () => {
    const response = await fetch('/api/chat', {
      method: 'POST',
      body: JSON.stringify({ messages: [{ role: 'user', content: 'Hello' }] }),
    });
 
    const reader = response.body!.getReader();
    const decoder = new TextDecoder();
    const tokens: string[] = [];
 
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      tokens.push(decoder.decode(value));
    }
 
    expect(tokens.length).toBeGreaterThan(1);
    expect(tokens.join('')).toContain('Hello');
  });
 
  it('should call tools when appropriate', async () => {
    const result = await processMessage('What is the weather in New York?');
    expect(result.toolCalls).toContainEqual(
      expect.objectContaining({ toolName: 'getWeather' })
    );
  });
 
  it('should maintain conversation context', async () => {
    const session = new ConversationSession();
    await session.send('My name is Alice');
    const response = await session.send('What is my name?');
    expect(response).toContain('Alice');
  });
});