/**
 * Gemini RAG Service
 *
 * Integrates with Google Gemini File Search API for RAG (Retrieval Augmented Generation).
 * Each user/organization gets their own FileSearchStore for document isolation.
 *
 * @see https://ai.google.dev/gemini-api/docs/file-search
 */

import { GoogleGenAI } from '@google/genai';
import { writeFileSync, unlinkSync, mkdtempSync, rmdirSync } from 'fs';
import { tmpdir } from 'os';
import { join } from 'path';
import { convexQuery, convexMutation } from '@/lib/convex-client';
import { logger } from '@/lib/logger';
import type {
  RAGQueryResponse,
  RAGCitation,
} from '@/lib/types/knowledge-base.types';

const scopedLogger = logger.scope('GeminiRAGService');

// Initialize Gemini client
let geminiClient: GoogleGenAI | null = null;

function getGeminiClient(): GoogleGenAI {
  if (!geminiClient) {
    const apiKey = process.env.GOOGLE_API_KEY;
    if (!apiKey) {
      throw new Error('GOOGLE_API_KEY environment variable is not set');
    }
    geminiClient = new GoogleGenAI({ apiKey });
  }
  return geminiClient;
}

// Default model for RAG queries
const RAG_MODEL = process.env.GEMINI_RAG_MODEL || 'gemini-2.5-flash';

// Maximum wait time for async operations (ms)
const MAX_WAIT_TIME = 300000; // 5 minutes
const POLL_INTERVAL = 3000;

/**
 * Helper to poll for async operation completion.
 * Note: uploadToFileSearchStore completes synchronously (SDK waits for upload),
 * so this is only needed for truly async operations like store creation.
 */
async function waitForOperation(
  client: GoogleGenAI,
  operation: { done?: boolean; name?: string }
): Promise<void> {
  if (operation.done) return;

  const startTime = Date.now();
  let currentOp = operation;

  while (!currentOp.done) {
    if (Date.now() - startTime > MAX_WAIT_TIME) {
      throw new Error('Operation timed out waiting for completion');
    }

    await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL));

    if (currentOp.name) {
      currentOp = await client.operations.get({
        operation: currentOp as Parameters<typeof client.operations.get>[0]['operation'],
      });
    } else {
      break;
    }
  }
}

export class GeminiRAGService {
  /**
   * Get or create a FileSearchStore for a user
   * Creates a new store on first document upload
   */
  static async getOrCreateStore(
    userId: string,
    displayName?: string
  ): Promise<{ storeId: string; storeName: string }> {
    const client = getGeminiClient();

    // Check if user already has a store
    const settings = await convexQuery<{
      _id: string;
      geminiFileSearchStoreId?: string | null;
      geminiFileSearchStoreName?: string | null;
    }>('voice_agents/queries:getKBSettingsByUserId', { userId });

    if (settings?.geminiFileSearchStoreId) {
      scopedLogger.debug('Using existing FileSearchStore', { userId });
      return {
        storeId: settings.geminiFileSearchStoreId,
        storeName: settings.geminiFileSearchStoreName || displayName || `kb-${userId}`,
      };
    }

    // Create a new FileSearchStore
    const storeName = displayName || `Knowledge Base - ${userId.slice(0, 8)}`;

    scopedLogger.info('Creating new FileSearchStore', { userId, storeName });

    try {
      const store = await client.fileSearchStores.create({
        config: { displayName: storeName },
      });

      if (!store.name) {
        throw new Error('Failed to create FileSearchStore: no name returned');
      }

      // Save store ID to user settings (upsert: query → update or create+update)
      const existing = settings ?? await convexQuery<{ _id: string }>(
        'voice_agents/queries:getKBSettingsByUserId', { userId }
      );

      if (existing?._id) {
        await convexMutation('voice_agents/mutations:updateKBSettings', {
          id: existing._id,
          geminiFileSearchStoreId: store.name,
          geminiFileSearchStoreName: storeName,
        });
      } else {
        // Create new settings, then update with Gemini fields
        const newId = await convexMutation<string>('voice_agents/mutations:createKBSettings', {
          userId,
          isEnabled: true,
          onboardingCompleted: false,
          maxChunksPerQuery: 5,
          similarityThreshold: 0.7,
        });
        if (newId) {
          await convexMutation('voice_agents/mutations:updateKBSettings', {
            id: newId,
            geminiFileSearchStoreId: store.name,
            geminiFileSearchStoreName: storeName,
          });
        }
      }

      scopedLogger.info('FileSearchStore created', { userId, storeId: store.name });

      return {
        storeId: store.name,
        storeName,
      };
    } catch (error) {
      scopedLogger.error('Failed to create FileSearchStore', error, { userId });
      throw error;
    }
  }

  /**
   * Upload a document to the user's FileSearchStore
   * Gemini handles chunking and embedding automatically
   */
  static async uploadDocument(
    userId: string,
    file: {
      buffer: Buffer;
      originalName: string;
      mimeType: string;
    },
    metadata?: {
      title?: string;
      category?: string;
    }
  ): Promise<{
    geminiFileName: string;
    geminiUploadedAt: Date;
  }> {
    const client = getGeminiClient();

    // Get or create the user's store
    const { storeId } = await this.getOrCreateStore(userId);

    scopedLogger.info('Uploading document to FileSearchStore', {
      userId,
      fileName: file.originalName,
      mimeType: file.mimeType,
      size: file.buffer.length,
    });

    // Create a temporary file for the upload (Node.js doesn't have File API)
    const tempDir = mkdtempSync(join(tmpdir(), 'gemini-upload-'));
    const tempFilePath = join(tempDir, file.originalName);

    try {
      // Write buffer to temp file
      writeFileSync(tempFilePath, file.buffer);

      scopedLogger.debug('Created temp file for upload', { tempFilePath });

      // Upload using file path (Gemini SDK accepts string path in Node.js)
      // uploadToFileSearchStore completes synchronously — SDK waits for upload to finish
      const operation = await client.fileSearchStores.uploadToFileSearchStore({
        fileSearchStoreName: storeId,
        file: tempFilePath,
        config: {
          displayName: metadata?.title || file.originalName,
          customMetadata: metadata?.category
            ? [{ key: 'category', stringValue: metadata.category }]
            : undefined,
        },
      });

      // Check for operation-level errors
      if (operation.error) {
        throw new Error(`Gemini upload failed: ${JSON.stringify(operation.error)}`);
      }

      // Extract the document name from the completed operation response
      const geminiFileName =
        operation.response?.documentName ||
        operation.name ||
        `uploaded-${Date.now()}`;
      const geminiUploadedAt = new Date();

      scopedLogger.info('Document uploaded successfully', {
        userId,
        geminiFileName,
        operationDone: operation.done,
        hasResponse: !!operation.response,
      });

      return {
        geminiFileName,
        geminiUploadedAt,
      };
    } catch (error) {
      scopedLogger.error('Failed to upload document', error, { userId });
      throw error;
    } finally {
      // Clean up temp file and directory
      try {
        unlinkSync(tempFilePath);
        rmdirSync(tempDir);
      } catch (cleanupError) {
        scopedLogger.warn('Failed to clean up temp file', { tempFilePath, error: cleanupError });
      }
    }
  }

  /**
   * Query the knowledge base using RAG
   * Uses Gemini FileSearch to find relevant documents and generate a response
   */
  static async queryWithRAG(
    userId: string,
    query: string,
    options?: {
      language?: 'en' | 'ar';
      maxResults?: number;
      includeCitations?: boolean;
      customInstructions?: string;
    }
  ): Promise<RAGQueryResponse> {
    const client = getGeminiClient();
    const startTime = Date.now();

    // Get user's FileSearchStore
    const settings = await convexQuery<{
      _id: string;
      geminiFileSearchStoreId?: string | null;
      customInstructions?: string | null;
      maxChunksPerQuery?: number | null;
    }>('voice_agents/queries:getKBSettingsByUserId', { userId });

    if (!settings?.geminiFileSearchStoreId) {
      return {
        success: false,
        answer: 'No knowledge base configured. Please upload documents first.',
        citations: [],
        metadata: {
          queryTimeMs: Date.now() - startTime,
          model: RAG_MODEL,
        },
      };
    }

    const language = options?.language || 'ar';
    const includeCitations = options?.includeCitations !== false;

    // Build the prompt
    let prompt = query;

    // Add language instruction
    if (language === 'ar') {
      prompt = `Please respond in Arabic. ${query}`;
    }

    // Add custom instructions if provided
    const instructions = options?.customInstructions || settings.customInstructions;
    if (instructions) {
      prompt = `${instructions}\n\nUser Query: ${prompt}`;
    }

    // Add citation instruction
    if (includeCitations) {
      prompt += '\n\nPlease cite the relevant sections from the documents when providing your answer.';
    }

    scopedLogger.info('Executing RAG query', {
      userId,
      queryLength: query.length,
      language,
    });

    try {
      const response = await client.models.generateContent({
        model: RAG_MODEL,
        contents: prompt,
        config: {
          tools: [
            {
              fileSearch: {
                fileSearchStoreNames: [settings.geminiFileSearchStoreId],
              },
            },
          ],
        },
      });

      // Extract citations from grounding metadata
      const citations: RAGCitation[] = [];
      const groundingChunks = response.candidates?.[0]?.groundingMetadata?.groundingChunks || [];

      for (const chunk of groundingChunks) {
        if (chunk.retrievedContext?.text) {
          citations.push({
            documentId: chunk.retrievedContext.uri || 'unknown',
            documentTitle: chunk.retrievedContext.title || 'Document',
            excerpt: chunk.retrievedContext.text.slice(0, 500),
          });
        }
      }

      const queryTimeMs = Date.now() - startTime;

      scopedLogger.info('RAG query completed', {
        userId,
        queryTimeMs,
        citationsCount: citations.length,
      });

      return {
        success: true,
        answer: response.text || 'No response generated.',
        citations: includeCitations ? citations : [],
        metadata: {
          queryTimeMs,
          model: RAG_MODEL,
        },
      };
    } catch (error) {
      scopedLogger.error('RAG query failed', error, { userId });

      return {
        success: false,
        answer: 'Failed to process your query. Please try again.',
        citations: [],
        metadata: {
          queryTimeMs: Date.now() - startTime,
          model: RAG_MODEL,
        },
      };
    }
  }

  /**
   * Delete a document from Gemini (cleanup when document is removed)
   * Note: Gemini FileSearch stores don't support individual file deletion,
   * so we mark it as deleted in our database and ignore in queries
   */
  static async markDocumentDeleted(
    userId: string,
    documentId: string
  ): Promise<void> {
    scopedLogger.info('Marking document as deleted', { userId, documentId });

    await convexMutation('voice_agents/mutations:updateKBDocument', {
      id: documentId,
      deletedAt: Date.now(),
    });
  }

  /**
   * Delete the entire FileSearchStore for a user
   * Used when account is deleted (after grace period)
   */
  static async deleteStore(userId: string): Promise<void> {
    const client = getGeminiClient();

    const settings = await convexQuery<{
      _id: string;
      geminiFileSearchStoreId?: string | null;
    }>('voice_agents/queries:getKBSettingsByUserId', { userId });

    if (!settings?.geminiFileSearchStoreId) {
      scopedLogger.warn('No FileSearchStore to delete', { userId });
      return;
    }

    scopedLogger.info('Deleting FileSearchStore', {
      userId,
      storeId: settings.geminiFileSearchStoreId,
    });

    try {
      await client.fileSearchStores.delete({
        name: settings.geminiFileSearchStoreId,
        config: { force: true },
      });

      // Clear store ID from settings
      await convexMutation('voice_agents/mutations:updateKBSettings', {
        id: settings._id,
        geminiFileSearchStoreId: '',
        geminiFileSearchStoreName: '',
      });

      scopedLogger.info('FileSearchStore deleted', { userId });
    } catch (error) {
      scopedLogger.error('Failed to delete FileSearchStore', error, { userId });
      throw error;
    }
  }

  /**
   * Get FileSearchStore statistics for a user
   */
  static async getStoreStats(userId: string): Promise<{
    hasStore: boolean;
    storeId: string | null;
    storeName: string | null;
    documentCount: number;
    activeDocumentCount: number;
  }> {
    const [settings, allDocs] = await Promise.all([
      convexQuery<{
        _id: string;
        geminiFileSearchStoreId?: string | null;
        geminiFileSearchStoreName?: string | null;
      }>('voice_agents/queries:getKBSettingsByUserId', { userId }),
      convexQuery<Array<{
        geminiFileName?: string | null;
        deletedAt?: number | null;
        status?: string;
      }>>('voice_agents/queries:getKBDocuments', { userId }),
    ]);

    const docs = allDocs ?? [];
    const totalDocs = docs.filter((d) => d.geminiFileName != null).length;
    const activeDocs = docs.filter(
      (d) => d.geminiFileName != null && d.deletedAt == null && d.status === 'COMPLETED'
    ).length;

    return {
      hasStore: !!settings?.geminiFileSearchStoreId,
      storeId: settings?.geminiFileSearchStoreId || null,
      storeName: settings?.geminiFileSearchStoreName || null,
      documentCount: totalDocs,
      activeDocumentCount: activeDocs,
    };
  }

  /**
   * Check if Gemini API is properly configured
   */
  static isConfigured(): boolean {
    return !!process.env.GOOGLE_API_KEY;
  }
}

// Export singleton-style functions for convenience
export const geminiRAGService = {
  getOrCreateStore: GeminiRAGService.getOrCreateStore.bind(GeminiRAGService),
  uploadDocument: GeminiRAGService.uploadDocument.bind(GeminiRAGService),
  queryWithRAG: GeminiRAGService.queryWithRAG.bind(GeminiRAGService),
  markDocumentDeleted: GeminiRAGService.markDocumentDeleted.bind(GeminiRAGService),
  deleteStore: GeminiRAGService.deleteStore.bind(GeminiRAGService),
  getStoreStats: GeminiRAGService.getStoreStats.bind(GeminiRAGService),
  isConfigured: GeminiRAGService.isConfigured.bind(GeminiRAGService),
};
