services/vectorService.js

import { ChromaClient } from 'chromadb';

class VectorService {
  constructor() {
    this.chromaClient = null;
    this.patientCollection = null;
    this.researchCollection = null;
    this.initialized = false;
  }
  
  async initialize() {
    if (this.initialized) return;
    
    try {
      this.chromaClient = new ChromaClient({
        path: './chroma_db'
      });
      
      this.patientCollection = await this._getOrCreateCollection('patient_records');
      this.researchCollection = await this._getOrCreateCollection('medical_research');
      this.initialized = true;
      console.log('[OK] Vector database initialized');
    } catch (error) {
      console.log(`[WARNING] Vector database initialization warning: ${error.message}`);
    }
  }
  
  async _getOrCreateCollection(name) {
    try {
      return await this.chromaClient.getCollection({ name });
    } catch {
      return await this.chromaClient.createCollection({
        name,
        metadata: { description: `Collection for ${name}` }
      });
    }
  }
  
  async addPatientRecord(patientId, patientData) {
    await this.initialize();
    
    const documents = [];
    const metadatas = [];
    const ids = [];
    
    // Add medical history
    if (patientData.medical_history) {
      patientData.medical_history.forEach((condition, idx) => {
        const docId = `patient_${patientId}_history_${idx}`;
        const conditionText = typeof condition === 'string' ? condition : condition.condition || 'Unknown';
        documents.push(`Patient ${patientData.name} has medical history: ${conditionText}`);
        metadatas.push({
          patient_id: patientId,
          patient_name: patientData.name,
          type: 'medical_history',
          condition: conditionText,
          age: String(patientData.age || ''),
          gender: patientData.gender || ''
        });
        ids.push(docId);
      });
    }
    
    // Add medications
    if (patientData.medications) {
      patientData.medications.forEach((med, idx) => {
        const docId = `patient_${patientId}_medication_${idx}`;
        let medText = `Patient ${patientData.name} is taking ${med.name}, dosage: ${med.dosage}, started: ${med.start_date}`;
        if (med.end_date) {
          medText += `, ended: ${med.end_date}`;
        }
        documents.push(medText);
        metadatas.push({
          patient_id: patientId,
          patient_name: patientData.name,
          type: 'medication',
          medication_name: med.name,
          dosage: med.dosage,
          start_date: med.start_date
        });
        ids.push(docId);
      });
    }
    
    // Add summary
    const medHistory = patientData.medical_history || ['None'];
    const medNames = (patientData.medications || []).map(m => m.name);
    const summaryText = `
      Patient Profile: ${patientData.name}, Age: ${patientData.age}, Gender: ${patientData.gender}.
      Medical History: ${medHistory.join(', ')}.
      Current Medications: ${medNames.join(', ') || 'None'}.
    `.trim();
    
    documents.push(summaryText);
    metadatas.push({
      patient_id: patientId,
      patient_name: patientData.name,
      type: 'summary',
      age: String(patientData.age || ''),
      gender: patientData.gender || ''
    });
    ids.push(`patient_${patientId}_summary`);
    
    if (documents.length > 0) {
      try {
        await this.patientCollection.upsert({
          documents,
          metadatas,
          ids
        });
        console.log(`   [OK] Indexed ${documents.length} records for patient ${patientData.name}`);
      } catch (error) {
        console.log(`   [ERROR] Error indexing patient ${patientData.name}: ${error.message}`);
      }
    }
    
    return documents.length;
  }
  
  async searchPatientRecords(query, patientId = null, nResults = 5) {
    await this.initialize();
    
    try {
      const whereFilter = patientId ? { patient_id: patientId } : undefined;
      
      const results = await this.patientCollection.query({
        queryTexts: [query],
        nResults,
        where: whereFilter
      });
      
      const formattedResults = [];
      if (results && results.documents && results.documents[0]) {
        for (let idx = 0; idx < results.documents[0].length; idx++) {
          formattedResults.push({
            document: results.documents[0][idx],
            metadata: results.metadatas[0][idx],
            distance: results.distances ? results.distances[0][idx] : null
          });
        }
      }
      
      return formattedResults;
    } catch (error) {
      console.log(`Search error: ${error.message}`);
      return [];
    }
  }
  
  async addResearchPaper(paperId, title, content, metadata) {
    await this.initialize();
    
    try {
      const chunkSize = 1000;
      const chunks = [];
      const words = content.split(' ');
      
      for (let i = 0; i < words.length; i += chunkSize) {
        const chunk = words.slice(i, i + chunkSize).join(' ');
        chunks.push(chunk);
      }
      
      const documents = [];
      const metadatas = [];
      const ids = [];
      
      chunks.forEach((chunk, idx) => {
        documents.push(chunk);
        const chunkMetadata = {
          paper_id: paperId,
          title,
          chunk_index: String(idx),
          total_chunks: String(chunks.length),
          ...Object.fromEntries(Object.entries(metadata).map(([k, v]) => [k, String(v)]))
        };
        metadatas.push(chunkMetadata);
        ids.push(`research_${paperId}_chunk_${idx}`);
      });
      
      if (documents.length > 0) {
        await this.researchCollection.upsert({
          documents,
          metadatas,
          ids
        });
      }
      
      return chunks.length;
    } catch (error) {
      console.log(`Error adding research paper: ${error.message}`);
      return 0;
    }
  }
  
  async searchResearchPapers(query, nResults = 5) {
    await this.initialize();
    
    try {
      const results = await this.researchCollection.query({
        queryTexts: [query],
        nResults
      });
      
      const formattedResults = [];
      if (results && results.documents && results.documents[0]) {
        for (let idx = 0; idx < results.documents[0].length; idx++) {
          formattedResults.push({
            document: results.documents[0][idx],
            metadata: results.metadatas[0][idx],
            distance: results.distances ? results.distances[0][idx] : null
          });
        }
      }
      
      return formattedResults;
    } catch (error) {
      console.log(`Research search error: ${error.message}`);
      return [];
    }
  }
  
  async deletePatientRecords(patientId) {
    await this.initialize();
    
    try {
      const results = await this.patientCollection.get({
        where: { patient_id: patientId }
      });
      
      if (results && results.ids && results.ids.length > 0) {
        await this.patientCollection.delete({ ids: results.ids });
      }
      return true;
    } catch (error) {
      console.log(`Error deleting patient records: ${error.message}`);
      return false;
    }
  }
  
  async getCollectionStats() {
    await this.initialize();
    
    try {
      const patientCount = await this.patientCollection.count();
      const researchCount = await this.researchCollection.count();
      
      return {
        patient_records_count: patientCount,
        research_papers_count: researchCount
      };
    } catch {
      return {
        patient_records_count: 0,
        research_papers_count: 0
      };
    }
  }
}

export const vectorService = new VectorService();