import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { HuggingFaceInferenceEmbeddings } from '@langchain/community/embeddings/hf';
import { ChatGroq } from '@langchain/groq';
import { Pinecone } from '@pinecone-database/pinecone';
import { PromptTemplate } from '@langchain/core/prompts';
import { StringOutputParser } from '@langchain/core/output_parsers';
import { RunnableSequence } from '@langchain/core/runnables';
import { formatDocumentsAsString } from 'langchain/util/document';
import NodeCache from 'node-cache';

@Injectable()
export class ChatbotService implements OnModuleInit {
  private readonly logger = new Logger(ChatbotService.name);
  private model: ChatGroq;
  private embeddings: HuggingFaceInferenceEmbeddings;
  private pineconeIndex: any;
  private retriever: any;
  private ragChain: any;
  private directLLMChain: any;
  private initialized = false;
  private useEmbeddings = false;

  private queryCache = new NodeCache({ stdTTL: 3600, checkperiod: 60 });
  private retrievalCache = new NodeCache({ stdTTL: 1800, checkperiod: 60 });

  constructor(private configService: ConfigService) {}

  // Initialize at module startup rather than on first request
  async onModuleInit(): Promise<void> {
    try {
      await this.initialize();
    } catch (error) {
      this.logger.error(`Failed to initialize service on startup: ${error.message}`);
      // Don't throw - let the application continue to start and retry initialization later
    }
  }

  public async initialize(): Promise<void> {
    if (this.initialized) return;

    try {
      const groqApiKey = this.configService.get<string>('GROQ_API_KEY');
      this.useEmbeddings = this.configService.get<string>('USE_EMBEDDINGS') === 'true';
      
      this.logger.log(`Initializing chatbot with embeddings: ${this.useEmbeddings}`);

      if (this.useEmbeddings) {
        const hfApiKey = this.configService.get<string>('HUGGINGFACE_HUB_API_KEY');
        const pineconeApiKey = this.configService.get<string>('PINECONE_API_KEY');
        const indexName = this.configService.get<string>('PINECONE_INDEX_NAME');

        // Validate RAG configs when embeddings are enabled
        this.validateRAGConfig({ groqApiKey, hfApiKey, pineconeApiKey, indexName });

        // Initialize services in parallel
        await Promise.all([
          this.initializeModel(groqApiKey),
          this.initializeEmbeddings(hfApiKey),
          this.initializePinecone(pineconeApiKey, indexName)
        ]);

        this.configureRAGChain();
      } else {
        // Only validate and initialize the LLM when not using embeddings
        this.validateDirectConfig({ groqApiKey });
        await this.initializeModel(groqApiKey);
        this.configureDirectLLMChain();
      }

      this.initialized = true;
      this.logger.log(`Chatbot service initialized successfully. Mode: ${this.useEmbeddings ? 'RAG with embeddings' : 'Direct LLM'}`);
    } catch (error) {
      this.logger.error(`Error initializing chatbot service: ${error.message}`);
      throw error;
    }
  }

  private validateRAGConfig(config: Record<string, string>): void {
    const missingConfigs = [];
    Object.entries(config).forEach(([key, value]) => {
      if (!value) missingConfigs.push(key);
    });

    if (missingConfigs.length > 0) {
      const errorMsg = `Missing RAG configuration: ${missingConfigs.join(', ')}`;
      this.logger.error(errorMsg);
      throw new Error(errorMsg);
    }
  }

  private validateDirectConfig(config: Record<string, string>): void {
    if (!config.groqApiKey) {
      const errorMsg = 'Missing configuration: GROQ_API_KEY';
      this.logger.error(errorMsg);
      throw new Error(errorMsg);
    }
  }

  private async initializeModel(apiKey: string): Promise<void> {
    try {
      this.model = new ChatGroq({
        apiKey,
        model: 'llama-3.1-8b-instant',
        temperature: 0.2,
      });
      this.logger.log('LLM model initialized successfully');
    } catch (error) {
      this.logger.error(`Failed to initialize LLM model: ${error.message}`);
      throw error;
    }
  }

  private async initializeEmbeddings(apiKey: string): Promise<void> {
    try {
      this.embeddings = new HuggingFaceInferenceEmbeddings({
        apiKey,
        model: 'sentence-transformers/all-mpnet-base-v2',
      });

      // Test connection
      await this.embeddings.embedQuery('Test connection');
      this.logger.log('Hugging Face embeddings initialized and tested successfully');
    } catch (error) {
      this.logger.error(`Hugging Face connection failed: ${error.message}`);
      throw new Error('Cannot proceed: Invalid Hugging Face API credentials');
    }
  }

  private async initializePinecone(apiKey: string, indexName: string): Promise<void> {
    try {
      const pinecone = new Pinecone({
        apiKey,
      });

      this.pineconeIndex = pinecone.Index(indexName);

      // Check index stats
      const stats = await this.pineconeIndex.describeIndexStats();
      this.logger.log(`Vector count in index: ${stats.totalVectorCount}`);
      if (stats.totalVectorCount === 0) {
        this.logger.warn('No vectors found in index! Make sure to load your data first.');
      }
    } catch (error) {
      this.logger.error(`Failed to initialize Pinecone: ${error.message}`);
      throw error;
    }
  }

  private configureRAGChain(): void {
    // Configure retriever with caching
    this.retriever = {
      invoke: async (query: string) => {
        try {
          // Check cache first
          const cachedDocuments = this.retrievalCache.get(query);
          if (cachedDocuments) {
            this.logger.log(`Retrieved documents from cache for query: "${query.substring(0, 30)}..."`);
            return cachedDocuments;
          }

          // Generate embedding
          const queryEmbedding = await this.embedQueryWithCache(query);
          this.logger.log(`Generated embedding for query: "${query.substring(0, 30)}..."`);

          // Query Pinecone
          const queryResponse = await this.pineconeIndex.query({
            vector: queryEmbedding,
            topK: 3,
            includeValues: false,
          });

          this.logger.log(`Retrieved ${queryResponse.matches.length} matches from Pinecone`);

          // Format documents
          const documents = queryResponse.matches.map(match => ({
            pageContent: match.metadata.pageContent || '',
            metadata: match.metadata,
          }));

          // Cache the results
          this.retrievalCache.set(query, documents);

          return documents;
        } catch (error) {
          this.logger.error(`Error in retriever: ${error.message}`);
          return [];
        }
      },
    };

    // const QA_PROMPT =
    //   PromptTemplate.fromTemplate(`You are a helpful assistant chatbot from CyberAlert Nepal that provides information based on the content from the provided documents.

    //   Context information is below:
    //   ---------------------
    //   {context}
    //   ---------------------
      
    //   Given this information, answer the question: {question}
      
    //   Important guidelines:
    //   1. Use information from the provided context to answer the question.
    //   2. If the context is empty or insufficient, respond in the same language format as the question:
    //      - In English: "I don't have enough information to answer this question."
    //      - In Romanized Nepali: "Ma sanga esprasna ko uttar din sapatta janakari xaina."
    //   3. Maintain the language format of the question:
    //      - If the question is in English, answer in English.
    //      - If the question is in Romanized Nepali, answer in Romanized Nepali.
    //   4. If the question is a general greeting, respond with a greeting in the same language format followed by "How can I help you?" or "Ma kasari madat garna sakchhu?" accordingly.
    //   5. If the question is about you, respond that you are a chatbot of CyberAlert Nepal and you can help with queries.
    //   6. Be precise, technically accurate, and informative.
    //   7. Do not make up or infer information not present in the context.
    //   8. Focus on directly answering the question.
    //   9. If you find relevant information in the context, make sure to use it in your answer.
    //   10. Always maintain a helpful and professional tone representing CyberAlert Nepal.
      
    //   Question: {question}
    //   Answer: `);

    const QA_PROMPT = 
    PromptTemplate.fromTemplate(`You are a helpful assistant chatbot from CyberAlert Nepal that provides information about cybersecurity and about CyberAlert Nepal.
  
    Answer the question: {question}
    
    Important guidelines:
    1. ONLY answer questions about cybersecurity topics or about CyberAlert Nepal.
    2. For any questions outside these two areas, politely decline in the same language format as the question:
       - In English: "I'm a specialized cybersecurity assistant for CyberAlert Nepal. I can only answer questions about cybersecurity or about our organization."
       - In Romanized Nepali: "Ma CyberAlertNepal ko chatbot ho. Ma cybersecurity ya hamro sanstha baare matra prashnaharu ko jawaf dina sakchhu."
    3. Maintain the language format of the question:
       - If the question is in English, answer in English.
       - If the question is in Romanized Nepali, answer in Romanized Nepali.
    4. If the question is a general greeting, respond with a greeting in the same language format followed by offering help with cybersecurity questions.
    5. When answering about CyberAlert Nepal, focus on the organization's services, mission, and cybersecurity offerings. CyberAlert Nepal provides cutting-edge cybersecurity solutions, digital forensics, and training. The organization is dedicated to ensuring a safe and secure online environment by combating cybercrime, digital fraud, and online threats through expert guidance and innovative technology.
    6. Be precise, technically accurate, and informative about cybersecurity topics.
    7. Do not make up or infer information that you're uncertain about.
    8. Always maintain a helpful and professional tone representing CyberAlert Nepal.
    
    Question: {question}
    Answer: `);
  
    this.ragChain = RunnableSequence.from([
      {
        context: async input => {
          const docs = await this.retriever.invoke(input.question);
          return formatDocumentsAsString(docs);
        },
        question: input => input.question,
      },
      QA_PROMPT,
      this.model,
      new StringOutputParser(),
    ]);

    this.logger.log('RAG chain configured');
  }

  private configureDirectLLMChain(): void {
    const DIRECT_PROMPT =
      PromptTemplate.fromTemplate(`You are a helpful assistant chatbot from CyberAlert Nepal.
      
      Important guidelines:
      1. Maintain the language format of the question:
         - If the question is in English, answer in English.
         - If the question is in Romanized Nepali, answer in Romanized Nepali.
      2. If the question is a general greeting, respond with a greeting in the same language format followed by "How can I help you?" or "Ma kasari madat garna sakchhu?" accordingly.
      3. If the question is about you, respond that you are a chatbot of CyberAlert Nepal and you can help with queries.
      4. Be precise, technically accurate, and informative.
      5. Focus on directly answering the question.
      6. Always maintain a helpful and professional tone representing CyberAlert Nepal.
      7. If you don't know the answer to a specific domain question, admit it politely.
      
      Question: {question}
      Answer: `);

    this.directLLMChain = RunnableSequence.from([
      {
        question: input => input.question,
      },
      DIRECT_PROMPT,
      this.model,
      new StringOutputParser(),
    ]);

    this.logger.log('Direct LLM chain configured');
  }

  private async embedQueryWithCache(query: string): Promise<any> {
    // Normalize query to improve cache hits (lowercase, trim whitespace)
    const normalizedQuery = query.toLowerCase().trim();

    const cachedEmbedding = this.queryCache.get(normalizedQuery);
    if (cachedEmbedding) {
      this.logger.log('Using cached embedding');
      return cachedEmbedding;
    }

    const embedding = await this.embeddings.embedQuery(normalizedQuery);
    this.queryCache.set(normalizedQuery, embedding);
    return embedding;
  }

  async processQuestion(
    question: string,
    chatHistory: Array<[string, string]> = [],
  ): Promise<{
    answer: string;
  }> {
    try {
      // Ensure service is initialized
      if (!this.initialized) {
        this.logger.log('Chatbot service not initialized. Initializing now...');
        await this.initialize();
      }

      if (!question || question.trim() === '') {
        throw new Error('Question cannot be empty');
      }

      this.logger.log(`Processing question: ${question}`);

      let answer = '';
      if (this.useEmbeddings) {
        // Generate answer using RAG chain with streaming
        this.logger.log('Using RAG chain with embeddings');
        const answerStream = await this.ragChain.stream({ question });
        for await (const chunk of answerStream) {
          answer += chunk;
        }
      } else {
        // Generate answer using direct LLM chain with streaming
        this.logger.log('Using direct LLM chain without embeddings');
        const answerStream = await this.directLLMChain.stream({ question });
        for await (const chunk of answerStream) {
          answer += chunk;
        }
      }

      return { answer };
    } catch (error) {
      this.logger.error(`Error processing question: ${error.message}`);
      throw new Error(`Failed to process your question: ${error.message}`);
    }
  }
}