import { BaseTool } from './BaseTool';
import { Vectorstore } from '../vectorstore/VectorstoreUpload';
import { Chunk } from '../types';
import * as fs from 'fs';
import { Networking } from '../../../../Network';

export class RAGTool extends BaseTool<{ hypothetical_document_chunk: string }> {
    constructor(
        private vectorstore: Vectorstore,
        summaries: () => string
    ) {
        super(
            'rag',
            'Perform a RAG search on user documents',
            {
                hypothetical_document_chunk: {
                    type: 'string',
                    description:
                        "Detailed version of the prompt that is effectively a hypothetical document chunk that would be ideal to embed and compare to the vectors of real document chunks to fetch the most relevant document chunks to answer the user's query",
                    required: 'true',
                },
            },
            `Your task is to first provide a response to the user's prompt based on the information given in the chunks and considering the chat history. Follow these steps:

            1. Carefully read and analyze the provided chunks, which may include text, images, or tables. Each chunk has an associated chunk_id.

            2. Review the prompt and chat history to understand the context of the user's question or request.

            3. Formulate a response that addresses the prompt using information from the relevant chunks. Your response should be informative and directly answer the user's question or request.

            4. Use citations to support your response. Citations should contain direct textual references to the granular, specific part of the original chunk that applies to the situation—with no text ommitted. Citations should be in the following format:
            - For text: <citation chunk_id="d980c2a7-cad3-4d7e-9eae-19bd2380bd02" type="text">relevant direct text from the chunk that the citation in referencing specifically</citation>
            - For images or tables: <citation chunk_id="9ef37681-b57e-4424-b877-e1ebc326ff11" type="image"></citation>

            Place citations after the sentences they apply to. You can use multiple citations in a row.

            5. If there's insufficient information in the provided chunks to answer the prompt sufficiently, ALWAYS respond with <answer>RAG not applicable</answer>

            Write your entire response, including follow-up questions, inside <answer> tags. Remember to use the citation format for both text and image references, and maintain a conversational tone throughout your response.

            !!!IMPORTANT Before you close the tag with </answer>, within the answer tags provide a set of 3 follow-up questions inside a <follow_up_questions> tag and individually within <question> tags. These should relate to the document, the current query, and the chat_history and should aim to help the user better understand whatever they are looking for.
            Also, ensure that the answer tags are wrapped with the correct step tags as well.`,

            `Performs a RAG (Retrieval-Augmented Generation) search on user documents and returns a 
            set of document chunks (either images or text) that can be used to provide a grounded response based on 
            user documents
            
            !!!IMPORTANT Use the RAG tool ANYTIME the question may potentially (even if you are not sure) relate to one of the user's documents.
            Here are the summaries of the user's documents:
            ${summaries()}`
        );
    }

    async execute(args: { hypothetical_document_chunk: string }): Promise<any> {
        const relevantChunks = await this.vectorstore.retrieve(args.hypothetical_document_chunk);
        const formatted_chunks = await this.getFormattedChunks(relevantChunks);
        return formatted_chunks;
    }

    async getFormattedChunks(relevantChunks: Chunk[]): Promise<{ type: string; text?: string; image_url?: { url: string } }[]> {
        try {
            const { formattedChunks } = await Networking.PostToServer('/formatChunks', { relevantChunks });

            if (!formattedChunks) {
                throw new Error('Failed to format chunks');
            }

            return formattedChunks;
        } catch (error) {
            console.error('Error formatting chunks:', error);
            throw error;
        }
    }
}