// SimilarityManager.js
// Manages cached embeddings and their staleness.
// Separated from EmbeddingManager for cleaner architecture

import { PriorityQueueFast } from '../core/PriorityQueueFast.js';
import { UnifiedKeyManager } from '../core/UnifiedKeyManager.js';

export class SimilarityManager {
  constructor(arbiter, options = {}) {
    this.arbiter = arbiter;
    
    // Initialize unified key manager
    this.keyManager = new UnifiedKeyManager();
    
    // JSON sketcher is an optional (deprecated) feature; inject via options to enable.
    this.jsonHasher = options.jsonHasher || null;

    // Cache converted vectors to avoid repeated conversions
    this.nodeKeyToVector = this.arbiter.cacheFactory(3000, {
      onEvict: (key, value) => {
        // Track evictions for debugging
        this.vectorCacheStats = this.vectorCacheStats || { evictions: 0 };
        this.vectorCacheStats.evictions++;
      },
      sampleSize: 5,
      sketchEpsilon: 0.01,
      sketchDelta: 0.01
    });
    
    // State related to the vector cache
    this.needsCacheRefresh = true; // Indicates if the nodeKeyToVector cache might be out of sync
    this.cachedVectorCount = 0;
    this.lastCacheRefreshTime = 0;
    
    // Warning reduction: Track missing embeddings to avoid spam
    this.missingEmbeddingNodes = new Set();
    this.embeddingWarningCount = 0;
    this.maxEmbeddingWarnings = 10; // Only show first 10 warnings
    
    // Performance tracking
    this.stats = {
      vectorCacheHits: 0,
      vectorConversions: 0,
      nodesWithoutEmbeddings: 0,
      staleEmbeddingsRefreshed: 0,
      cacheRefreshes: 0,
    };

    // Staleness Management for Embeddings
    this.staleNodesForEmbeddingQueue = new PriorityQueueFast();
    this.staleNodesIndex = new Map(); // Key: nodeKey, Value: StaleNodeItem
  }

  /**
   * Convert string key to integer for PriorityQueueFast
   * @private
   */
  _stringToInt(key) {
    return this.keyManager.getStringId(key);
  }

  /**
   * Convert integer key back to string for PriorityQueueFast
   * @private
   */
  _intToString(intKey) {
    return this.keyManager.getIdString(intKey);
  }

  /**
   * Rebuild the entire nodeKeyToVector cache from scratch.
   * This is useful if there's a global change or to ensure consistency.
   */
  refreshAllCachedVectors() {
    console.log('🔄 Refreshing all cached vectors...');
    const startTime = process.hrtime.bigint();
    
    this.nodeKeyToVector.clear();
    let addedCount = 0;
    let skippedCount = 0;
    
    for (const [nodeId, node] of this.arbiter.nodes.entries()) {
      if (node.embedding && node.key) {
          const vector = this._convertEmbeddingToVector(node.embedding, node.key);
        if (vector) {
          this.nodeKeyToVector.set(node.key, vector);
          addedCount++;
        } else {
          skippedCount++;
        }
      } else {
        skippedCount++;
      }
    }
    
    const refreshTime = Number(process.hrtime.bigint() - startTime) / 1_000_000;
    console.log(`✅ Cached vectors refreshed in ${refreshTime.toFixed(2)}ms`);
    console.log(`   Added ${addedCount} vectors, skipped ${skippedCount}`);
    
    this.needsCacheRefresh = false;
    this.cachedVectorCount = addedCount;
    this.lastCacheRefreshTime = Date.now();
    this.stats.cacheRefreshes++;
    
    // Clear the stale queue as all embeddings are now considered fresh from the cache perspective
    this.staleNodesForEmbeddingQueue.clear();
    this.staleNodesIndex.clear();
    
    return addedCount > 0;
  }

  /**
   * Convert embedding to vector format with caching and statistics
   */
  _convertEmbeddingToVector(embedding, nodeKey) {
    this.stats.vectorConversions++;
    
    if (embedding.sketch) {
      return new Uint32Array(embedding.sketch);
    } else if (embedding instanceof Uint32Array) {
      return embedding; // Keep as Uint32Array
    } else {
      console.warn(`Unknown embedding format for node ${nodeKey}`);
      return null;
    }
  }

  /**
   * Ensure the cache is ready.
   */
  ensureCacheReady() {
    if (this.needsCacheRefresh) {
      return this.refreshAllCachedVectors();
    }
    return true;
  }

  /**
   * Check if a node's vector is cached.
   */
  isNodeVectorCached(nodeKey) {
    return this.nodeKeyToVector.has(nodeKey);
  }

  /**
   * Mark cache for refresh (e.g. when many embeddings change)
   */
  markCacheForRefresh() {
    this.needsCacheRefresh = true;
  }

  /**
   * Add/Update a single node's vector in the cache.
   */
  updateNodeVector(nodeKey, embedding) {
    if (!embedding) {
      this.nodeKeyToVector.delete(nodeKey); // Remove if no embedding
      // console.warn(`[SimilarityManager] No embedding provided for node ${nodeKey}, removing from vector cache if present.`);
      return false;
    }
    
    const vector = this._convertEmbeddingToVector(embedding, nodeKey);
    if (vector) {
      this.nodeKeyToVector.set(nodeKey, vector);
      this.cachedVectorCount = this.nodeKeyToVector.size;
      return true;
    } else {
      // console.warn(`[SimilarityManager] Failed to convert embedding for node ${nodeKey}, not updating vector cache.`);
      return false;
    }
  }

  /**
   * Remove a node's vector from the cache.
   */
  removeNodeVector(nodeKey) {
    if (this.nodeKeyToVector.has(nodeKey)) {
      this.nodeKeyToVector.delete(nodeKey);
      this.cachedVectorCount = this.nodeKeyToVector.size;
        return true;
    }
    return false;
  }

  /**
   * Get comprehensive statistics including all optimizations
   */
  getStats() {
    return {
      ...this.stats,
      cachedVectors: this.cachedVectorCount,
      needsCacheRefresh: this.needsCacheRefresh,
      cacheReady: !this.needsCacheRefresh,
      lastCacheRefreshTime: this.lastCacheRefreshTime,
      timeSinceCacheRefresh: Date.now() - this.lastCacheRefreshTime,
      vectorCacheSize: this.nodeKeyToVector.size,
      vectorCacheHitRate: this.stats.vectorCacheHits / Math.max(this.stats.vectorConversions, 1),
      missingEmbeddingCount: this.missingEmbeddingNodes.size,
      embeddingWarningsSuppressed: Math.max(0, this.stats.nodesWithoutEmbeddings - this.maxEmbeddingWarnings),
      estimatedMemoryUsage: {
        vectorCache: this.nodeKeyToVector.size * (256 * 4), // Assuming Uint32Array, 256 elements, 4 bytes each
      }
    };
  }

  /**
   * Get embedding summary for diagnostics
   */
  getEmbeddingSummary() {
    const totalNodesWithAttemptedCache = this.missingEmbeddingNodes.size + this.cachedVectorCount;
    return {
      totalNodesCheckedForCache: totalNodesWithAttemptedCache,
      nodesWithCachedVectors: this.cachedVectorCount,
      nodesWithoutEmbeddingsForCache: this.missingEmbeddingNodes.size,
      vectorCacheCoverage: this.cachedVectorCount / Math.max(totalNodesWithAttemptedCache, 1),
      warningsShown: Math.min(this.embeddingWarningCount, this.maxEmbeddingWarnings),
      warningsSuppressed: Math.max(0, this.stats.nodesWithoutEmbeddings - this.maxEmbeddingWarnings),
      sampleMissingNodes: Array.from(this.missingEmbeddingNodes).slice(0, 5)
    };
  }

  /**
   * Get debug information for a specific node's cached vector
   */
  getNodeDebugInfo(nodeKey) {
    const nodeId = this.arbiter.nodeIdByKey.get(nodeKey);
    const node = nodeId !== undefined ? this.arbiter.nodes.get(nodeId) : null;
    const cachedVector = this.nodeKeyToVector.get(nodeKey);
    
    let hasSketch = false;
    let embeddingDimension = undefined;
    
    if (node?.embedding) {
      if (node.embedding.sketch) {
        hasSketch = true;
        embeddingDimension = node.embedding.sketch.length;
      } else if (node.embedding instanceof Uint32Array) {
        hasSketch = false; // Or true if this is considered the "sketch"
        embeddingDimension = node.embedding.length; 
      }
    }
    
    return {
      nodeKey,
      nodeId,
      exists: node !== null,
      hasEmbedding: node?.embedding !== null && node?.embedding !== undefined,
      hasSketch, // Indicates if original embedding had a .sketch property
      isVectorCached: cachedVector !== undefined,
      cachedVectorDimension: cachedVector?.length,
      rawEmbeddingDimension: embeddingDimension,
      nodeType: node?.type
    };
  }

  /**
   * Clear all state including optimization caches (for testing)
   */
  clear() {
    this.nodeKeyToVector.clear();
    
    if (this.jsonHasher && typeof this.jsonHasher.clearNodeStringCache === 'function') {
      this.jsonHasher.clearNodeStringCache();
    }

    this.needsCacheRefresh = true;
    this.cachedVectorCount = 0;
    this.lastCacheRefreshTime = 0;
    
    this.stats = {
      vectorCacheHits: 0,
      vectorConversions: 0,
      nodesWithoutEmbeddings: 0,
      staleEmbeddingsRefreshed: 0,
      cacheRefreshes: 0,
    };
    this.missingEmbeddingNodes.clear();
    this.embeddingWarningCount = 0;
    this.staleNodesForEmbeddingQueue.clear();
    this.staleNodesIndex.clear();
  }

  /**
   * Marks a node as stale, requiring its embedding to be refreshed and re-indexed.
   * @param {string | number} nodeKeyOrId - The key or ID of the node.
   * @param {string} reason - Why the node is stale (e.g., 'node_data_changed', 'relation_change').
   * @param {number} accessTime - Optional timestamp for last access, defaults to now.
   */
  markNodeAsStaleForEmbedding(nodeKeyOrId, reason, accessTime = Date.now()) {
    const nodeKey = (typeof nodeKeyOrId === 'string') 
      ? nodeKeyOrId 
      : this.arbiter.nodeManager.getNodeKey(nodeKeyOrId);

    if (!nodeKey) {
      // console.warn(`[SimilarityManager] Cannot mark unknown node as stale: ${nodeKeyOrId}`);
      return;
    }

    const now = Date.now();
    if (this.staleNodesIndex.has(nodeKey)) {
      const existingItem = this.staleNodesIndex.get(nodeKey);
      existingItem.lastAccessTime = Math.max(existingItem.lastAccessTime, accessTime); // Take the latest access time
      existingItem.staleSinceTime = Math.min(existingItem.staleSinceTime, now); // Take the earliest stale time
      existingItem.reason = reason; // Update reason if new one provided
      // If using a heap-based PQ that needs explicit update:
      // this.staleNodesForEmbeddingQueue.update(existingItem);
    } else {
      const staleItem = {
        id: nodeKey, // Use nodeKey as the ID for the PQ item
        itemType: 'nodeEmbedding', // Could be 'cachedVector' now
        lastAccessTime: accessTime,
        staleSinceTime: now,
        reason: reason
      };
      // Use nodeKey as integer key and staleSinceTime as priority (lower = higher priority)
      const nodeKeyInt = this._stringToInt(nodeKey);
      this.staleNodesForEmbeddingQueue.push(nodeKeyInt, now);
      this.staleNodesIndex.set(nodeKey, staleItem);
    }
    // this.needsRebuild = true; // Now relates to the vector cache
    this.needsCacheRefresh = true; // Mark the main vector cache for refresh if any node is stale
  }

  /**
   * Refreshes embeddings for a specified number of stale nodes.
   * This involves getting the fresh embedding via EmbeddingManager and updating the local vector cache.
   * @param {number} countToRefresh - The maximum number of stale nodes to refresh.
   * @returns {number} The number of items still remaining in the stale queue.
   */
  refreshStaleEmbeddings(countToRefresh) {
    let refreshedCount = 0;
    for (let i = 0; i < countToRefresh; i++) {
      if (this.staleNodesForEmbeddingQueue.isEmpty()) {
        break;
      }
      const nodeKeyInt = this.staleNodesForEmbeddingQueue.pop();
      const nodeKey = this._intToString(nodeKeyInt);
      const staleItem = this.staleNodesIndex.get(nodeKey);
      this.staleNodesIndex.delete(nodeKey);

      if (!staleItem) {
        console.warn(`[SimilarityManager] Stale item not found for key ${nodeKey} during refresh.`);
        continue;
      }
      const nodeId = this.arbiter.nodeIdByKey.get(nodeKey);

      if (nodeId === undefined) {
        console.warn(`[SimilarityManager] NodeId not found for key ${nodeKey} during stale refresh.`);
        continue;
      }

      // 1. Ensure the node's embedding is up-to-date via EmbeddingManager
      if (this.arbiter.embeddingManager && typeof this.arbiter.embeddingManager.ensureFreshEmbedding === 'function') {
        this.arbiter.embeddingManager.ensureFreshEmbedding(nodeId);
      } else {
        console.warn(`[SimilarityManager] EmbeddingManager or ensureFreshEmbedding not available for node ${nodeKey}.`);
        // Decide if we should proceed or skip. For now, try to proceed with current embedding.
      }

      // 2. Retrieve the (potentially updated) embedding from the node object
      const node = this.arbiter.nodes.get(nodeId);
      if (!node) {
        console.warn(`[SimilarityManager] Node object not found for id ${nodeId} (key: ${nodeKey}) during stale refresh.`);
        continue;
      }

      // 3. Update this manager's vector cache
      if (node.embedding) {
        this.updateNodeVector(nodeKey, node.embedding);
          // console.log(`[SimilarityManager] Refreshed vector cache for node ${nodeKey} (Reason: ${staleItem.reason}).`);
      } else {
        // console.warn(`[SimilarityManager] Node ${nodeKey} still has no embedding after refresh attempt.`);
        this.removeNodeVector(nodeKey); // Remove from cache if no current embedding
      }
      
      refreshedCount++;
      this.stats.staleEmbeddingsRefreshed++;
    }

    if (refreshedCount > 0) {
      // If specific items were refreshed, the main cache might not need a full rebuild,
      // but individual items are now fresh.
      // For simplicity, one could still mark for full refresh, or rely on individual updates.
      // this.needsCacheRefresh = true; // Or manage more granularly.
    }
    return this.staleNodesForEmbeddingQueue.size();
  }

  /**
   * Rebuild the similarity index after mass embedding updates
   * This is a placeholder for future similarity index optimization
   */
  rebuildIndex() {
    // Currently the similarity index is built on-demand through the nodeKeyToVector cache
    // This method is here for compatibility with EmbeddingManager
    // In the future, we could add more sophisticated indexing here
    this.needsCacheRefresh = true;
    return true; // Return success indicator for tests
  }
} 
