import { pipeline, env } from "@xenova/transformers";
import { QdrantClient } from "@qdrant/js-client-rest";
import crypto from "crypto";

env.allowRemoteModels = true;
env.localModelPath = "./models";

export class QdrantMemoryService {
  private client: QdrantClient;
  private collectionName = "kabelfachmann";
  private embedder: any = null;

  constructor(
    url: string = process.env.QDRANT_URL || "http://qdrant-mcp:6333",
  ) {
    this.client = new QdrantClient({ url });
  }

  async initialize() {
    console.error("Loading embedding model...");
    this.embedder = await pipeline(
      "feature-extraction",
      "Xenova/all-MiniLM-L6-v2",
    );

    console.error(`Checking for collection: ${this.collectionName}`);
    try {
      const collections = await this.client.getCollections();
      const exists = collections.collections.some(
        (c) => c.name === this.collectionName,
      );

      if (!exists) {
        console.error(`Creating collection: ${this.collectionName}`);
        await this.client.createCollection(this.collectionName, {
          vectors: {
            size: 384,
            distance: "Cosine",
          },
        });
        console.error("Collection created successfully.");
      }
    } catch (e) {
      console.error("Failed to initialize Qdrant collection:", e);
      throw e;
    }
  }

  private async getEmbedding(text: string): Promise<number[]> {
    if (!this.embedder) {
      throw new Error("Embedder not initialized. Call initialize() first.");
    }
    const output = await this.embedder(text, {
      pooling: "mean",
      normalize: true,
    });
    return Array.from(output.data);
  }

  async storeMemory(label: string, content: string): Promise<boolean> {
    try {
      const vector = await this.getEmbedding(content);
      const id = crypto.randomUUID();

      await this.client.upsert(this.collectionName, {
        wait: true,
        points: [
          {
            id,
            vector,
            payload: { label, content, timestamp: new Date().toISOString() },
          },
        ],
      });
      return true;
    } catch (e) {
      console.error("Failed to store memory:", e);
      return false;
    }
  }

  async retrieveMemory(
    query: string,
    limit: number = 5,
  ): Promise<Array<{ label: string; content: string; score: number }>> {
    try {
      const vector = await this.getEmbedding(query);
      const searchResults = await this.client.search(this.collectionName, {
        vector,
        limit,
        with_payload: true,
      });

      return searchResults.map((result) => ({
        label: String(result.payload?.label || ""),
        content: String(result.payload?.content || ""),
        score: result.score,
      }));
    } catch (e) {
      console.error("Failed to retrieve memory:", e);
      return [];
    }
  }
}