import { Injectable } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import * as fs from 'fs';
import * as path from 'path';
import matter from 'gray-matter';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
import { HuggingFaceInferenceEmbeddings } from '@langchain/community/embeddings/hf';
import { Document } from 'langchain/document';
import { Pinecone } from '@pinecone-database/pinecone';

@Injectable()
export class MdxService {
  private embeddings: HuggingFaceInferenceEmbeddings;
  private pinecone: Pinecone;
  private readonly indexName: string;

  constructor(private configService: ConfigService) {

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

    if (!hfApiKey) {
      throw new Error('HUGGINGFACE_HUB_API_KEY not found in configuration');
    }
    if (!pineconeApiKey) {
      throw new Error('PINECONE_API_KEY not found in configuration');
    }
    if (!this.indexName) {
      throw new Error('PINECONE_INDEX_NAME not found in configuration');
    }

    this.embeddings = new HuggingFaceInferenceEmbeddings({
      apiKey: hfApiKey,
      model: 'sentence-transformers/all-mpnet-base-v2',
    });

    this.pinecone = new Pinecone({
      apiKey: pineconeApiKey,
    });

  }

  private async testHFConnection(): Promise<boolean> {
    try {
      await this.embeddings.embedQuery('Test connection');
      return true;
    } catch (error) {
      console.error(`Hugging Face connection failed: ${error.message}`);
      return false;
    }
  }

  async processMdxFiles(directoryPath: string): Promise<void> {

    const connectionValid = await this.testHFConnection();
    if (!connectionValid) {
      throw new Error('Cannot proceed: Invalid Mistral API credentials');
    }

    if (!fs.existsSync(directoryPath)) {
      console.error(`Directory not found: ${directoryPath}`);
      throw new Error(`Directory not found: ${directoryPath}`);
    }

    const mdxFiles = this.getAllMdxFiles(directoryPath);

    if (mdxFiles.length === 0) {
      return;
    }

    const documents: Document[] = [];

    for (const filePath of mdxFiles) {
      try {
        const { content, metadata } = await this.parseMdxFile(filePath);
        const textChunks = await this.splitTextIntoChunks(content);

        for (const [index, chunk] of textChunks.entries()) {
          documents.push(
            new Document({
              pageContent: chunk,
              metadata: {
                ...metadata,
                source: filePath,
                chunk: index,
              },
            }),
          );
        }
      } catch (error) {
        console.error(`Error processing file ${filePath}: ${error.message}`);
        console.error(error.stack);
      }
    }

    if (documents.length > 0) {
      await this.storeDocumentsInPinecone(documents);
    } else {
    }
  }

  private getAllMdxFiles(directoryPath: string): string[] {
    const allFiles: string[] = [];

    const readDirectory = (dirPath: string) => {
      try {
        const entries = fs.readdirSync(dirPath, { withFileTypes: true });

        for (const entry of entries) {
          const fullPath = path.join(dirPath, entry.name);

          if (entry.isDirectory()) {
            readDirectory(fullPath);
          } else if (entry.isFile() && (entry.name.endsWith('.mdx') || entry.name.endsWith('.md'))) {
            allFiles.push(fullPath);
          }
        }
      } catch (error) {
        console.error(`Error reading directory ${dirPath}: ${error.message}`);
      }
    };

    readDirectory(directoryPath);
    return allFiles;
  }

  private async parseMdxFile(filePath: string): Promise<{ content: string; metadata: any }> {
    const fileContent = fs.readFileSync(filePath, 'utf-8');
    const { content, data } = matter(fileContent);

    try {
      const withoutCodeBlocks = content.replace(/``````/g, '');

      const withoutHtml = withoutCodeBlocks.replace(/<[^>]*>/g, ' ');

      const withoutImports = withoutHtml.replace(/import[\s\S]*?from\s+(['"]).*?\1;?/g, '');

      const withoutExports = withoutImports.replace(/export[\s\S]*?;/g, '');

      const plainText = withoutExports
        .replace(/#{1,6}\s+(.*)/g, '$1') // Replace headers with their content
        .replace(/\*\*(.*?)\*\*/g, '$1') // Remove bold
        .replace(/\*(.*?)\*/g, '$1') // Remove italic
        .replace(/\[(.*?)\]\(.*?\)/g, '$1') // Replace links with their text
        .replace(/\s+/g, ' ') // Replace multiple spaces with single space
        .trim();


      return {
        content: plainText,
        metadata: {
          ...data,
          filePath,
          fileName: path.basename(filePath),
        },
      };
    } catch (error) {
      console.error(`Error parsing MDX file ${filePath}: ${error.message}`);
      throw error;
    }
  }

  private async splitTextIntoChunks(text: string): Promise<string[]> {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 1000,
      chunkOverlap: 200,
    });

    const chunks = await splitter.splitText(text);
    return chunks;
  }

  private async storeDocumentsInPinecone(documents: Document[]): Promise<void> {
    try {

      const pineconeApiKey = this.configService.get<string>('PINECONE_API_KEY');
      if (!pineconeApiKey) {
        throw new Error('PINECONE_API_KEY not found or empty in configuration');
      }

      const pinecone = new Pinecone({
        apiKey: pineconeApiKey,
      });

      await this.ensureIndexExists(pinecone);

      const index = pinecone.Index(this.indexName);

      const batchSize = 100;
      const totalBatches = Math.ceil(documents.length / batchSize);


      let successCount = 0;
      const concurrencyLimit = 3;

      for (let i = 0; i < documents.length; i += batchSize * concurrencyLimit) {
        const batchPromises = [];

        for (let j = 0; j < concurrencyLimit; j++) {
          const startIdx = i + j * batchSize;
          if (startIdx < documents.length) {
            const batchNumber = Math.floor(startIdx / batchSize) + 1;
            const batch = documents.slice(startIdx, startIdx + batchSize);
            batchPromises.push(this.processBatch(batch, batchNumber, totalBatches, index));
          }
        }

        const results = await Promise.all(batchPromises);
        successCount += results.reduce((sum, count) => sum + count, 0);
      }

    } catch (error) {
      console.error(`Error in document storage process: ${error.message}`);
      console.error(error.stack);
      throw error;
    }
  }

  private async ensureIndexExists(pinecone: Pinecone): Promise<void> {
    const indexList = await pinecone.listIndexes();
    const indexNames = indexList.indexes?.map(index => index.name) || [];

    if (!indexNames.includes(this.indexName)) {
      await pinecone.createIndex({
        name: this.indexName,
        dimension: 768,
        metric: 'cosine',
        spec: {
          serverless: {
            cloud: 'aws',
            region: 'us-east-1',
          },
        },
      });

      let isReady = false;
      const maxAttempts = 15;
      let attempts = 0;

      while (!isReady && attempts < maxAttempts) {
        attempts++;
        await new Promise(resolve => setTimeout(resolve, 5000));

        try {
          const indexDescription = await pinecone.describeIndex(this.indexName);
          if (indexDescription.status?.ready) {
            isReady = true;
          } else {
          }
        } catch (error) {
        }
      }

      if (!isReady) {
        console.warn(`Index may not be fully ready, but proceeding after ${attempts * 5} seconds`);
      }
    } else {
    }
  }

  private async processBatch(batch: Document[], batchNumber: number, totalBatches: number, index: any): Promise<number> {

    const vectors = await this.generateEmbeddingsForBatch(batch);

    if (vectors.length === 0) {
      return 0;
    }

    const maxRetries = 3;
    let retryCount = 0;
    let success = false;
    let successCount = 0;

    while (!success && retryCount < maxRetries) {
      try {
        await index.upsert(vectors);
        success = true;
        successCount = vectors.length;
      } catch (error) {
        retryCount++;
        const delay = Math.min(2000 * Math.pow(2, retryCount), 10000); // Exponential backoff with max 10s
        console.error(`Error uploading vectors (attempt ${retryCount}/${maxRetries}): ${error.message}`);

        if (retryCount >= maxRetries) {
          successCount = await this.performIndividualUploads(vectors, index);
        } else {
          await new Promise(resolve => setTimeout(resolve, delay));
        }
      }
    }

    return successCount;
  }

  private async generateEmbeddingsForBatch(batch: Document[]): Promise<any[]> {
    const vectors = [];

    const embeddingPromises = batch.map(async doc => {
      try {
        const embedding = await this.embeddings.embedQuery(doc.pageContent);

        return {
          id: `doc_${doc.metadata.fileName}_${doc.metadata.chunk}`,
          values: embedding,
          metadata: {
            ...doc.metadata,
            pageContent: doc.pageContent,
          },
        };
      } catch (error) {
        console.error(`Error generating embedding: ${error.message}`);
        return null;
      }
    });

    const results = await Promise.all(embeddingPromises);
    return results.filter(result => result !== null);
  }

  private async performIndividualUploads(vectors: any[], index: any): Promise<number> {
    let successCount = 0;

    for (const vector of vectors) {
      try {
        await index.upsert([vector]);
        successCount++;

        await new Promise(resolve => setTimeout(resolve, 100));
      } catch (error) {
        console.error(`Failed to upload vector ${vector.id}: ${error.message}`);
      }
    }

    return successCount;
  }
}