// ============================================================
// Persistent Agent Memory
// Based on ReasoningBank (Google Research, ICLR 2026)
//
// 3-phase pipeline:
// 1. select_memory() - embedding retrieval from history
// 2. evaluate_trajectory() - assess success/failure
// 3. induce_memory() - LLM extracts lessons from experiences
//
// P3 Priority from TMLPD Roadmap
// ============================================================

export interface MemoryBlock {
  id: string;
  type: 'experience' | 'lesson' | 'pattern' | 'context';
  content: string;
  embedding: number[];
  agentId: string;
  taskType?: string;
  success: boolean;
  quality: number;           // 0-1 quality of this memory
  useCount: number;          // How many times this was retrieved
  createdAt: number;
  lastUsedAt: number;
  expiresAt?: number;        // Optional TTL
}

export interface Trajectory {
  id: string;
  agentId: string;
  taskId: string;
  query: string;
  steps: TrajectoryStep[];
  finalResult: string;
  success: boolean;
  costUSD: number;
  durationMs: number;
  quality: number;            // 0-1 self-reported or judged quality
  createdAt: number;
}

export interface TrajectoryStep {
  step: number;
  action: string;
  observation: string;
  thought: string;
  success: boolean;
}

export interface MemoryConfig {
  maxMemories: number;       // Max memories per agent
  maxAgeDays: number;        // Days before memory expires
  similarityThreshold: number; // Min similarity to retrieve
  embeddingModel: string;
  inductionEnabled: boolean;  // Enable LLM-based memory induction
}

const DEFAULT_CONFIG: Required<MemoryConfig> = {
  maxMemories: 1000,
  maxAgeDays: 30,
  similarityThreshold: 0.7,
  embeddingModel: 'text-embedding-3-small',
  inductionEnabled: true,
};

// ============================================================
// Simple Embedding (cosine similarity)
// ============================================================

function cosineSimilarity(a: number[], b: number[]): number {
  if (a.length !== b.length) return 0;
  
  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) + 1e-8);
}

function generateSimpleEmbedding(text: string): number[] {
  // Simple hash-based pseudo-embedding for demo
  // In production, use real embeddings from OpenAI/Cohere
  const dim = 384;
  const embedding = new Array(dim).fill(0);
  
  for (let i = 0; i < text.length; i++) {
    const char = text.charCodeAt(i);
    embedding[i % dim] += char * (i + 1);
    embedding[(i * 31) % dim] += char;
  }
  
  // Normalize
  const norm = Math.sqrt(embedding.reduce((sum, v) => sum + v * v, 0));
  return embedding.map(v => v / (norm + 1e-8));
}

// ============================================================
// Persistent Memory Store
// ============================================================

export class PersistentMemory {
  private memories: Map<string, MemoryBlock[]> = new Map();
  private trajectories: Map<string, Trajectory[]> = new Map();
  private config: Required<MemoryConfig>;

  constructor(config: Partial<MemoryConfig> = {}) {
    this.config = { ...DEFAULT_CONFIG, ...config };
  }

  /**
   * Get memories for an agent
   */
  getMemories(agentId: string): MemoryBlock[] {
    return this.memories.get(agentId) || [];
  }

  /**
   * Select relevant memories using embedding similarity
   */
  selectMemory(
    agentId: string,
    query: string,
    options?: {
      limit?: number;
      taskType?: string;
      minSimilarity?: number;
    }
  ): MemoryBlock[] {
    const memories = this.getMemories(agentId);
    const queryEmbedding = generateSimpleEmbedding(query);
    const limit = options?.limit || 5;
    const threshold = options?.minSimilarity || this.config.similarityThreshold;
    
    // Score and filter memories
    const scored = memories
      .map(memory => ({
        memory,
        similarity: cosineSimilarity(queryEmbedding, memory.embedding),
      }))
      .filter(item => item.similarity >= threshold)
      .filter(item => {
        if (options?.taskType && item.memory.taskType !== options.taskType) {
          return false;
        }
        // Check expiry
        if (item.memory.expiresAt && item.memory.expiresAt < Date.now()) {
          return false;
        }
        return true;
      })
      .sort((a, b) => {
        // Prioritize: similarity * quality * recency
        const scoreA = a.similarity * a.memory.quality * (1 + a.memory.useCount * 0.1);
        const scoreB = b.similarity * b.memory.quality * (1 + b.memory.useCount * 0.1);
        return scoreB - scoreA;
      })
      .slice(0, limit);
    
    // Update use count
    for (const item of scored) {
      item.memory.useCount++;
      item.memory.lastUsedAt = Date.now();
    }
    
    return scored.map(item => item.memory);
  }

  /**
   * Evaluate a trajectory and score it
   */
  evaluateTrajectory(trajectory: Trajectory): {
    quality: number;
    lessons: string[];
    patterns: string[];
  } {
    const lessons: string[] = [];
    const patterns: string[] = [];

    // Quality based on success and efficiency
    let quality = 0.5;
    
    if (trajectory.success) {
      quality += 0.3;
    }
    
    // Efficiency bonus (fewer steps = better)
    if (trajectory.steps.length <= 3) {
      quality += 0.1;
    } else if (trajectory.steps.length > 10) {
      quality -= 0.1;
    }
    
    // Cost efficiency
    if (trajectory.costUSD < 0.01) {
      quality += 0.1;
    }
    
    // Extract patterns from steps
    const failedSteps = trajectory.steps.filter(s => !s.success);
    if (failedSteps.length > 0) {
      patterns.push(`had_${failedSteps.length}_failures`);
      lessons.push(`Recovery needed from ${failedSteps.length} step failures`);
    }
    
    const longSteps = trajectory.steps.filter(s => s.observation.length > 500);
    if (longSteps.length > 0) {
      patterns.push('verbose_observations');
    }
    
    // Check for retry patterns
    const retryCount = trajectory.steps.filter(s => 
      s.action.includes('retry') || s.action.includes('again')
    ).length;
    if (retryCount > 2) {
      patterns.push('multi_retry');
      lessons.push('Consider different approach instead of retry');
    }
    
    // Check for successful tool sequences
    const toolSequence = trajectory.steps
      .filter(s => s.success && s.action)
      .map(s => s.action)
      .join(' → ');
    if (toolSequence) {
      patterns.push(`tool_flow:${toolSequence}`);
    }
    
    return {
      quality: Math.min(1, Math.max(0, quality)),
      lessons,
      patterns,
    };
  }

  /**
   * Induce memory from successful trajectory
   * Uses LLM to extract key lessons (simplified version)
   */
  induceMemory(trajectory: Trajectory): MemoryBlock[] {
    const { quality, lessons, patterns } = this.evaluateTrajectory(trajectory);
    const newMemories: MemoryBlock[] = [];
    
    // Only store memories from decent quality experiences
    if (quality < 0.3) {
      return newMemories;
    }
    
    // 1. Experience memory
    if (trajectory.success && trajectory.finalResult) {
      const experienceMemory: MemoryBlock = {
        id: `exp_${trajectory.id}`,
        type: 'experience',
        content: `Task: ${trajectory.query}\nResult: ${trajectory.finalResult.slice(0, 500)}`,
        embedding: generateSimpleEmbedding(trajectory.query + ' ' + trajectory.finalResult),
        agentId: trajectory.agentId,
        taskType: this.classifyTask(trajectory.query),
        success: true,
        quality,
        useCount: 0,
        createdAt: Date.now(),
        lastUsedAt: Date.now(),
        expiresAt: Date.now() + this.config.maxAgeDays * 24 * 60 * 60 * 1000,
      };
      newMemories.push(experienceMemory);
    }
    
    // 2. Lesson memories
    for (const lesson of lessons) {
      const lessonMemory: MemoryBlock = {
        id: `lesson_${trajectory.id}_${newMemories.length}`,
        type: 'lesson',
        content: lesson,
        embedding: generateSimpleEmbedding(lesson),
        agentId: trajectory.agentId,
        taskType: this.classifyTask(trajectory.query),
        success: trajectory.success,
        quality: quality * 0.8,
        useCount: 0,
        createdAt: Date.now(),
        lastUsedAt: Date.now(),
        expiresAt: Date.now() + this.config.maxAgeDays * 2 * 24 * 60 * 60 * 1000,
      };
      newMemories.push(lessonMemory);
    }
    
    // 3. Pattern memories
    for (const pattern of patterns) {
      if (pattern.startsWith('tool_flow:')) {
        const patternMemory: MemoryBlock = {
          id: `pattern_${trajectory.id}_${newMemories.length}`,
          type: 'pattern',
          content: `Effective tool sequence: ${pattern.replace('tool_flow:', '')}`,
          embedding: generateSimpleEmbedding(pattern),
          agentId: trajectory.agentId,
          taskType: this.classifyTask(trajectory.query),
          success: trajectory.success,
          quality: quality * 0.9,
          useCount: 0,
          createdAt: Date.now(),
          lastUsedAt: Date.now(),
        };
        newMemories.push(patternMemory);
      }
    }
    
    return newMemories;
  }

  /**
   * Store memories for an agent
   */
  storeMemories(agentId: string, newMemories: MemoryBlock[]): void {
    const existing = this.getMemories(agentId);
    
    // Add new memories
    existing.push(...newMemories);
    
    // Prune old memories if over limit
    if (existing.length > this.config.maxMemories) {
      // Sort by quality * recency and keep top N
      const sorted = existing
        .sort((a, b) => {
          const scoreA = a.quality * (1 + (Date.now() - a.lastUsedAt) / (1000 * 60 * 60 * 24));
          const scoreB = b.quality * (1 + (Date.now() - b.lastUsedAt) / (1000 * 60 * 60 * 24));
          return scoreB - scoreA;
        })
        .slice(0, this.config.maxMemories);
      
      this.memories.set(agentId, sorted);
    } else {
      this.memories.set(agentId, existing);
    }
  }

  /**
   * Store a trajectory and induce memories
   */
  storeTrajectory(trajectory: Trajectory): void {
    // Store trajectory
    if (!this.trajectories.has(trajectory.agentId)) {
      this.trajectories.set(trajectory.agentId, []);
    }
    this.trajectories.get(trajectory.agentId)!.push(trajectory);
    
    // Induce and store memories if enabled
    if (this.config.inductionEnabled) {
      const newMemories = this.induceMemory(trajectory);
      this.storeMemories(trajectory.agentId, newMemories);
    }
  }

  /**
   * Get trajectories for an agent
   */
  getTrajectories(agentId: string, limit?: number): Trajectory[] {
    const trajectories = this.trajectories.get(agentId) || [];
    const sorted = trajectories.sort((a, b) => b.createdAt - a.createdAt);
    return limit ? sorted.slice(0, limit) : sorted;
  }

  /**
   * Build context with retrieved memories
   */
  buildContext(
    agentId: string,
    query: string,
    options?: { maxLength?: number }
  ): string {
    const memories = this.selectMemory(agentId, query, { limit: 5 });
    const maxLength = options?.maxLength || 2000;
    
    if (memories.length === 0) {
      return '';
    }
    
    const parts = ['[Relevant Past Experience]'];
    
    for (const memory of memories) {
      const prefix = memory.type === 'lesson' ? '💡 ' : 
                     memory.type === 'pattern' ? '🔄 ' : '📝 ';
      const line = `${prefix}${memory.content.slice(0, 300)}`;
      parts.push(line);
    }
    
    const context = parts.join('\n');
    return context.length > maxLength ? context.slice(0, maxLength) + '...' : context;
  }

  /**
   * Classify task type from query
   */
  private classifyTask(query: string): string {
    const lower = query.toLowerCase();
    
    if (/\b(code|function|debug|implement)\b/.test(lower)) return 'code';
    if (/\b(why|how|analyze|compare)\b/.test(lower)) return 'reasoning';
    if (/\b(write|story|creative|brainstorm)\b/.test(lower)) return 'creative';
    if (/\b(deploy|production|security)\b/.test(lower)) return 'critical';
    
    return 'general';
  }

  /**
   * Get memory statistics
   */
  getStats(agentId: string): {
    totalMemories: number;
    byType: Record<string, number>;
    avgQuality: number;
    totalTrajectories: number;
  } {
    const memories = this.getMemories(agentId);
    const trajectories = this.getTrajectories(agentId);
    
    const byType: Record<string, number> = {};
    let totalQuality = 0;
    
    for (const memory of memories) {
      byType[memory.type] = (byType[memory.type] || 0) + 1;
      totalQuality += memory.quality;
    }
    
    return {
      totalMemories: memories.length,
      byType,
      avgQuality: memories.length > 0 ? totalQuality / memories.length : 0,
      totalTrajectories: trajectories.length,
    };
  }

  /**
   * Clear all memories for an agent
   */
  clear(agentId: string): void {
    this.memories.delete(agentId);
    this.trajectories.delete(agentId);
  }

  /**
   * Export memories as JSON
   */
  export(agentId: string): { memories: MemoryBlock[]; trajectories: Trajectory[] } {
    return {
      memories: this.getMemories(agentId),
      trajectories: this.getTrajectories(agentId),
    };
  }

  /**
   * Import memories from JSON
   */
  import(agentId: string, data: { memories?: MemoryBlock[]; trajectories?: Trajectory[] }): void {
    if (data.memories) {
      this.memories.set(agentId, data.memories);
    }
    if (data.trajectories) {
      this.trajectories.set(agentId, data.trajectories);
    }
  }
}

// ============================================================
// Convenience export
// ============================================================

export const persistentMemory = new PersistentMemory();