import dotenv from 'dotenv';
import { XMLBuilder, XMLParser } from 'fast-xml-parser';
import { escape } from 'lodash'; // Imported escape from lodash
import OpenAI from 'openai';
import { DocumentOptions } from '../../../../documents/Documents';
import { AnswerParser } from '../response_parsers/AnswerParser';
import { StreamedAnswerParser } from '../response_parsers/StreamedAnswerParser';
import { BaseTool } from '../tools/BaseTool';
import { CalculateTool } from '../tools/CalculateTool';
import { CreateDocTool } from '../tools/CreateDocumentTool';
import { DataAnalysisTool } from '../tools/DataAnalysisTool';
import { NoTool } from '../tools/NoTool';
import { SearchTool } from '../tools/SearchTool';
import { Parameter, ParametersType, TypeMap } from '../types/tool_types';
import { AgentMessage, ASSISTANT_ROLE, AssistantMessage, Observation, PROCESSING_TYPE, ProcessingInfo, TEXT_TYPE } from '../types/types';
import { Vectorstore } from '../vectorstore/Vectorstore';
import { getReactPrompt } from './prompts';
import { ChatCompletionMessageParam } from 'openai/resources';
import { Doc } from '../../../../../fields/Doc';
import { parsedDoc } from '../chatboxcomponents/ChatBox';
import { WebsiteInfoScraperTool } from '../tools/WebsiteInfoScraperTool';
import { Upload } from '../../../../../server/SharedMediaTypes';
import { RAGTool } from '../tools/RAGTool';
import { GPTTutorialTool } from '../tools/TutorialTool';

dotenv.config();

/**
 * The Agent class handles the interaction between the assistant and the tools available,
 * processes user queries, and manages the communication flow between the tools and OpenAI.
 */
export class Agent {
    // Private properties
    private client: OpenAI;
    private messages: AgentMessage[] = [];
    private interMessages: AgentMessage[] = [];
    private vectorstore: Vectorstore;
    private _history: () => string;
    private _summaries: () => string;
    private _csvData: () => { filename: string; id: string; text: string }[];
    private actionNumber: number = 0;
    private thoughtNumber: number = 0;
    private processingNumber: number = 0;
    private processingInfo: ProcessingInfo[] = [];
    private streamedAnswerParser: StreamedAnswerParser = new StreamedAnswerParser();
    private tools: Record<string, BaseTool<ReadonlyArray<Parameter>>>;
    private Document: Doc;

    /**
     * The constructor initializes the agent with the vector store and toolset, and sets up the OpenAI client.
     * @param _vectorstore Vector store instance for document storage and retrieval.
     * @param summaries A function to retrieve document summaries.
     * @param history A function to retrieve chat history.
     * @param csvData A function to retrieve CSV data linked to the assistant.
     * @param addLinkedUrlDoc A function to add a linked document from a URL.
     * @param createCSVInDash A function to create a CSV document in the dashboard.
     */
    constructor(
        _vectorstore: Vectorstore,
        summaries: () => string,
        history: () => string,
        csvData: () => { filename: string; id: string; text: string }[],
        addLinkedUrlDoc: (url: string, id: string) => void,
        createImage: (result: Upload.FileInformation & Upload.InspectionResults, options: DocumentOptions) => void,
        addLinkedDoc: (doc: parsedDoc) => Doc | undefined,
        createCSVInDash: (url: string, title: string, id: string, data: string) => void,
        document: Doc
    ) {
        // Initialize OpenAI client with API key from environment
        this.client = new OpenAI({ apiKey: process.env.OPENAI_KEY, dangerouslyAllowBrowser: true });
        this.vectorstore = _vectorstore;
        this._history = history;
        this._summaries = summaries;
        this._csvData = csvData;
        this.Document = document;

        // Define available tools for the assistant
        this.tools = {
            calculate: new CalculateTool(),
            rag: new RAGTool(this.vectorstore),
            dataAnalysis: new DataAnalysisTool(csvData),
            websiteInfoScraper: new WebsiteInfoScraperTool(addLinkedUrlDoc),
            searchTool: new SearchTool(addLinkedUrlDoc),
            noTool: new NoTool(),
            createDoc: new CreateDocTool(addLinkedDoc),
            generateTutorialNode: new GPTTutorialTool(addLinkedDoc),
        };
    }

    /**
     * This method handles the conversation flow with the assistant, processes user queries,
     * and manages the assistant's decision-making process, including tool actions.
     * @param question The user's question.
     * @param onProcessingUpdate Callback function for processing updates.
     * @param onAnswerUpdate Callback function for answer updates.
     * @param maxTurns The maximum number of turns to allow in the conversation.
     * @returns The final response from the assistant.
     */
    async askAgent(question: string, onProcessingUpdate: (processingUpdate: ProcessingInfo[]) => void, onAnswerUpdate: (answerUpdate: string) => void, maxTurns: number = 30): Promise<AssistantMessage> {
        console.log(`Starting query: ${question}`);
        const MAX_QUERY_LENGTH = 1000; // adjust the limit as needed

        // Check if the question exceeds the maximum length
        if (question.length > MAX_QUERY_LENGTH) {
            return { role: ASSISTANT_ROLE.ASSISTANT, content: [{ text: 'User query too long. Please shorten your question and try again.', index: 0, type: TEXT_TYPE.NORMAL, citation_ids: null }], processing_info: [] };
        }

        const sanitizedQuestion = escape(question); // Sanitized user input

        // Push sanitized user's question to message history
        this.messages.push({ role: 'user', content: sanitizedQuestion });

        // Retrieve chat history and generate system prompt
        const chatHistory = this._history();
        let systemPrompt = getReactPrompt(Object.values(this.tools), this._summaries, chatHistory);

        // If this is a Dash documentation assistant chat, modify the system prompt
        if (this.Document?.is_dash_doc_assistant) {
            systemPrompt = systemPrompt.replace(
                '<task>',
                `<task>
                IMPORTANT: You are specifically focused on helping users with questions about Dash documentation and usage. When users ask questions, interpret them in the context of Dash documentation and features, even if they don't explicitly mention Dash. For example, if a user asks "How do I create a document?", interpret this as "How do I create a document in Dash?" and provide relevant Dash-specific guidance.
                
                For any questions about Dash features, functionality, or usage, you should use the generateTutorialNode tool to create a tutorial document that explains the concept in detail. This tool will help create well-formatted, interactive tutorials that guide users through Dash features.`
            );
        }

        // Initialize intermediate messages
        this.interMessages = [{ role: 'system', content: systemPrompt }];

        this.interMessages.push({
            role: 'user',
            content: this.constructUserPrompt(1, 'user', `<query>${sanitizedQuestion}</query>`),
        });

        // Setup XML parser and builder
        const parser = new XMLParser({
            ignoreAttributes: false,
            attributeNamePrefix: '@_',
            textNodeName: '_text',
            isArray: name => name === 'url',
            processEntities: false, // Disable processing of entities
            stopNodes: ['*.entity'], // Do not process any entities
        });
        const builder = new XMLBuilder({ ignoreAttributes: false, attributeNamePrefix: '@_' });

        let currentAction: string | undefined;
        this.processingInfo = [];

        let i = 2;
        while (i < maxTurns) {
            console.log(this.interMessages);
            console.log(`Turn ${i}/${maxTurns}`);

            // eslint-disable-next-line no-await-in-loop
            const result = await this.execute(onProcessingUpdate, onAnswerUpdate);
            this.interMessages.push({ role: 'assistant', content: result });

            i += 2;

            let parsedResult;
            try {
                // Parse XML result from the assistant
                parsedResult = parser.parse(result);

                // Validate the structure of the parsedResult
                this.validateAssistantResponse(parsedResult);
            } catch (error) {
                throw new Error(`Error parsing or validating response: ${error}`);
            }

            // Extract the stage from the parsed result
            const stage = parsedResult.stage;
            if (!stage) {
                throw new Error(`Error: No stage found in response`);
            }

            // Handle different stage elements (thoughts, actions, inputs, answers)
            for (const key in stage) {
                if (key === 'thought') {
                    // Handle assistant's thoughts
                    console.log(`Thought: ${stage[key]}`);
                    this.processingNumber++;
                } else if (key === 'action') {
                    // Handle action stage
                    currentAction = stage[key] as string;
                    console.log(`Action: ${currentAction}`);

                    if (this.tools[currentAction]) {
                        // Prepare the next action based on the current tool
                        const nextPrompt = [
                            {
                                type: 'text',
                                text: `<stage number="${i + 1}" role="user">` + builder.build({ action_rules: this.tools[currentAction].getActionRule() }) + `</stage>`,
                            } as Observation,
                        ];
                        this.interMessages.push({ role: 'user', content: nextPrompt });
                        break;
                    } else {
                        // Handle error in case of an invalid action
                        console.log('Error: No valid action');
                        this.interMessages.push({
                            role: 'user',
                            content: `<stage number="${i + 1}" role="system-error-reporter">No valid action, try again.</stage>`,
                        });
                        break;
                    }
                } else if (key === 'action_input') {
                    // Handle action input stage
                    const actionInput = stage[key];
                    console.log(`Action input full:`, actionInput);
                    console.log(`Action input:`, actionInput.inputs);

                    if (currentAction) {
                        try {
                            // Process the action with its input
                            // eslint-disable-next-line no-await-in-loop
                            const observation = (await this.processAction(currentAction, actionInput.inputs)) as Observation[];
                            const nextPrompt = [{ type: 'text', text: `<stage number="${i + 1}" role="user"> <observation>` }, ...observation, { type: 'text', text: '</observation></stage>' }] as Observation[];
                            console.log(observation);
                            this.interMessages.push({ role: 'user', content: nextPrompt });
                            this.processingNumber++;
                            break;
                        } catch (error) {
                            throw new Error(`Error processing action: ${error}`);
                        }
                    } else {
                        throw new Error('Error: Action input without a valid action');
                    }
                } else if (key === 'answer') {
                    // If an answer is found, end the query
                    console.log('Answer found. Ending query.');
                    this.streamedAnswerParser.reset();
                    const parsedAnswer = AnswerParser.parse(result, this.processingInfo);
                    return parsedAnswer;
                }
            }
        }

        throw new Error('Reached maximum turns. Ending query.');
    }

    private constructUserPrompt(stageNumber: number, role: string, content: string): string {
        return `<stage number="${stageNumber}" role="${role}">${content}</stage>`;
    }

    /**
     * Executes a step in the conversation, processing the assistant's response and parsing it in real-time.
     * @param onProcessingUpdate Callback for processing updates.
     * @param onAnswerUpdate Callback for answer updates.
     * @returns The full response from the assistant.
     */
    private async execute(onProcessingUpdate: (processingUpdate: ProcessingInfo[]) => void, onAnswerUpdate: (answerUpdate: string) => void): Promise<string> {
        // Stream OpenAI response for real-time updates
        const stream = await this.client.chat.completions.create({
            model: 'gpt-4o',
            messages: this.interMessages as ChatCompletionMessageParam[],
            temperature: 0,
            stream: true,
            stop: ['</stage>'],
        });

        let fullResponse: string = '';
        let currentTag: string = '';
        let currentContent: string = '';
        let isInsideTag: boolean = false;

        // Process each chunk of the streamed response
        for await (const chunk of stream) {
            const content = chunk.choices[0]?.delta?.content || '';
            fullResponse += content;

            // Parse the streamed content character by character
            for (const char of content) {
                if (currentTag === 'answer') {
                    // Handle answer parsing for real-time updates
                    currentContent += char;
                    const streamedAnswer = this.streamedAnswerParser.parse(char);
                    onAnswerUpdate(streamedAnswer);
                    continue;
                } else if (char === '<') {
                    // Start of a new tag
                    isInsideTag = true;
                    currentTag = '';
                    currentContent = '';
                } else if (char === '>') {
                    // End of the tag
                    isInsideTag = false;
                    if (currentTag.startsWith('/')) {
                        currentTag = '';
                    }
                } else if (isInsideTag) {
                    // Append characters to the tag name
                    currentTag += char;
                } else if (currentTag === 'thought' || currentTag === 'action_input_description') {
                    // Handle processing information for thought or action input description
                    currentContent += char;
                    const current_info = this.processingInfo.find(info => info.index === this.processingNumber);
                    if (current_info) {
                        current_info.content = currentContent.trim();
                        onProcessingUpdate(this.processingInfo);
                    } else {
                        this.processingInfo.push({
                            index: this.processingNumber,
                            type: currentTag === 'thought' ? PROCESSING_TYPE.THOUGHT : PROCESSING_TYPE.ACTION,
                            content: currentContent.trim(),
                        });
                        onProcessingUpdate(this.processingInfo);
                    }
                }
            }
        }

        return fullResponse;
    }

    /**
     * Validates the assistant's response to ensure it conforms to the expected XML structure.
     * @param response The parsed XML response from the assistant.
     * @throws An error if the response does not meet the expected structure.
     */
    private validateAssistantResponse(response: { stage: { [key: string]: object | string } }) {
        if (!response.stage) {
            throw new Error('Response does not contain a <stage> element');
        }

        // Validate that the stage has the required attributes
        const stage = response.stage;
        if (!stage['@_number'] || !stage['@_role']) {
            throw new Error('Stage element must have "number" and "role" attributes');
        }

        // Extract the role of the stage to determine expected content
        const role = stage['@_role'];

        // Depending on the role, validate the presence of required elements
        if (role === 'assistant') {
            // Assistant's response should contain either 'thought', 'action', 'action_input', or 'answer'
            if (!('thought' in stage || 'action' in stage || 'action_input' in stage || 'answer' in stage)) {
                throw new Error('Assistant stage must contain a thought, action, action_input, or answer element');
            }

            // If 'thought' is present, validate it
            if ('thought' in stage) {
                if (typeof stage.thought !== 'string' || stage.thought.trim() === '') {
                    throw new Error('Thought must be a non-empty string');
                }
            }

            // If 'action' is present, validate it
            if ('action' in stage) {
                if (typeof stage.action !== 'string' || stage.action.trim() === '') {
                    throw new Error('Action must be a non-empty string');
                }

                // Optional: Check if the action is among allowed actions
                const allowedActions = Object.keys(this.tools);
                if (!allowedActions.includes(stage.action)) {
                    throw new Error(`Action "${stage.action}" is not a valid tool`);
                }
            }

            // If 'action_input' is present, validate its structure
            if ('action_input' in stage) {
                const actionInput = stage.action_input as object;

                if (!('action_input_description' in actionInput) || typeof actionInput.action_input_description !== 'string') {
                    throw new Error('action_input must contain an action_input_description string');
                }

                if (!('inputs' in actionInput)) {
                    throw new Error('action_input must contain an inputs object');
                }

                // Further validation of inputs can be done here based on the expected parameters of the action
            }

            // If 'answer' is present, validate its structure
            if ('answer' in stage) {
                const answer = stage.answer as object;

                // Ensure answer contains at least one of the required elements
                if (!('grounded_text' in answer || 'normal_text' in answer)) {
                    throw new Error('Answer must contain grounded_text or normal_text');
                }

                // Validate follow_up_questions
                if (!('follow_up_questions' in answer)) {
                    throw new Error('Answer must contain follow_up_questions');
                }

                // Validate loop_summary
                if (!('loop_summary' in answer)) {
                    throw new Error('Answer must contain a loop_summary');
                }

                // Additional validation for citations, grounded_text, etc., can be added here
            }
        } else if (role === 'user') {
            // User's stage should contain 'query' or 'observation'
            if (!('query' in stage || 'observation' in stage)) {
                throw new Error('User stage must contain a query or observation element');
            }

            // Validate 'query' if present
            if ('query' in stage && typeof stage.query !== 'string') {
                throw new Error('Query must be a string');
            }

            // Validate 'observation' if present
            if ('observation' in stage) {
                // Ensure observation has the correct structure
                // This can be expanded based on how observations are structured
            }
        } else {
            throw new Error(`Unknown role "${role}" in stage`);
        }

        // Add any additional validation rules as necessary
    }

    /**
     * Helper function to check if a string can be parsed as an array of the expected type.
     * @param input The input string to check.
     * @param expectedType The expected type of the array elements ('string', 'number', or 'boolean').
     * @returns The parsed array if valid, otherwise throws an error.
     */
    private parseArray<T>(input: string, expectedType: 'string' | 'number' | 'boolean'): T[] {
        try {
            // Parse the input string into a JSON object
            const parsed = JSON.parse(input);

            // Check if the parsed object is an array and if all elements are of the expected type
            if (Array.isArray(parsed) && parsed.every(item => typeof item === expectedType)) {
                return parsed;
            } else {
                throw new Error(`Invalid ${expectedType} array format.`);
            }
        } catch (error) {
            throw new Error(`Failed to parse ${expectedType} array: ` + error);
        }
    }

    /**
     * Processes a specific action by invoking the appropriate tool with the provided inputs.
     * This method ensures that the action exists and validates the types of `actionInput`
     * based on the tool's parameter rules. It throws errors for missing required parameters
     * or mismatched types before safely executing the tool with the validated input.
     *
     * NOTE: In the future, it should typecheck for specific tool parameter types using the `TypeMap` or otherwise.
     *
     * Type validation includes checks for:
     * - `string`, `number`, `boolean`
     * - `string[]`, `number[]` (arrays of strings or numbers)
     *
     * @param action The action to perform. It corresponds to a registered tool.
     * @param actionInput The inputs for the action, passed as an object where each key is a parameter name.
     * @returns A promise that resolves to an array of `Observation` objects representing the result of the action.
     * @throws An error if the action is unknown, if required parameters are missing, or if input types don't match the expected parameter types.
     */
    private async processAction(action: string, actionInput: ParametersType<ReadonlyArray<Parameter>>): Promise<Observation[]> {
        // Check if the action exists in the tools list
        if (!(action in this.tools)) {
            throw new Error(`Unknown action: ${action}`);
        }
        console.log(actionInput);

        for (const param of this.tools[action].parameterRules) {
            // Check if the parameter is required and missing in the input
            if (param.required && !(param.name in actionInput) && !this.tools[action].inputValidator(actionInput)) {
                throw new Error(`Missing required parameter: ${param.name}`);
            }

            // Check if the parameter type matches the expected type
            const expectedType = param.type.replace('[]', '') as 'string' | 'number' | 'boolean';
            const isArray = param.type.endsWith('[]');
            const input = actionInput[param.name];

            if (isArray) {
                // Check if the input is a valid array of the expected type
                const parsedArray = this.parseArray(input as string, expectedType);
                actionInput[param.name] = parsedArray as TypeMap[typeof param.type];
            } else if (input !== undefined && typeof input !== expectedType) {
                throw new Error(`Invalid type for parameter ${param.name}: expected ${expectedType}`);
            }
        }

        const tool = this.tools[action];

        return await tool.execute(actionInput);
    }
}