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

// Be sure to set local caching options for transformers
env.allowRemoteModels = true;
env.localModelPath = './models';

export class QdrantMemoryService {
    private client: QdrantClient;
    private collectionName = 'mcp_memory';
    private embedder: any = null;

    constructor(url: string = 'http://localhost:6333') {
        this.client = new QdrantClient({ url });
    }

    /**
     * Initializes the embedding model and the Qdrant collection
     */
    async initialize() {
        // 1. Load the embedding model (using a lightweight model suitable for semantic search)
        console.error('Loading embedding model...');
        this.embedder = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');

        // 2. Ensure collection exists
        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, // size for all-MiniLM-L6-v2
                        distance: 'Cosine'
                    }
                });
                console.error('Collection created successfully.');
            }
        } catch (e) {
            console.error('Failed to initialize Qdrant collection:', e);
            throw e;
        }
    }

    /**
     * Generates a vector embedding for the given text
     */
    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);
    }

    /**
     * Stores a memory entry into Qdrant
     */
    async storeMemory(label: string, content: string): Promise<boolean> {
        try {
            const fullText = `${label}: ${content}`;
            const vector = await this.getEmbedding(fullText);
            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;
        }
    }

    /**
     * Retrieves memory entries relevant to the query
     */
    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 [];
        }
    }
}