import OpenAI from 'openai';
import { Tool, AgentMessage, AssistantMessage, TEXT_TYPE, CHUNK_TYPE, ASSISTANT_ROLE, ProcessingInfo, PROCESSING_TYPE } from './types';
import { getReactPrompt } from './prompts';
import { XMLParser, XMLBuilder } from 'fast-xml-parser';
import { Vectorstore } from './vectorstore/Vectorstore';
import { ChatCompletionMessageParam } from 'openai/resources';
import dotenv from 'dotenv';
import { CalculateTool } from './tools/CalculateTool';
import { RAGTool } from './tools/RAGTool';
import { DataAnalysisTool } from './tools/DataAnalysisTool';
import { WebsiteInfoScraperTool } from './tools/WebsiteInfoScraperTool';
import { SearchTool } from './tools/SearchTool';
import { NoTool } from './tools/NoTool';
import { on } from 'events';
import { v4 as uuidv4 } from 'uuid';
import { AnswerParser } from './AnswerParser';
import { StreamedAnswerParser } from './StreamedAnswerParser';

dotenv.config();

export class Agent {
    private client: OpenAI;
    private tools: Record<string, Tool<any>>;
    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();

    constructor(_vectorstore: Vectorstore, summaries: () => string, history: () => string, csvData: () => { filename: string; id: string; text: string }[], addLinkedUrlDoc: (url: string, id: string) => void) {
        this.client = new OpenAI({ apiKey: process.env.OPENAI_KEY, dangerouslyAllowBrowser: true });
        this.vectorstore = _vectorstore;
        this._history = history;
        this._summaries = summaries;
        this._csvData = csvData;
        this.tools = {
            calculate: new CalculateTool(),
            rag: new RAGTool(this.vectorstore),
            dataAnalysis: new DataAnalysisTool(csvData),
            websiteInfoScraper: new WebsiteInfoScraperTool(addLinkedUrlDoc),
            searchTool: new SearchTool(addLinkedUrlDoc),
            no_tool: new NoTool(),
        };
    }

    async askAgent(question: string, onProcessingUpdate: (processingUpdate: ProcessingInfo[]) => void, onAnswerUpdate: (answerUpdate: string) => void, maxTurns: number = 30): Promise<AssistantMessage> {
        console.log(`Starting query: ${question}`);
        this.messages.push({ role: 'user', content: question });
        const chatHistory = this._history();
        const systemPrompt = getReactPrompt(Object.values(this.tools), this._summaries, chatHistory);
        this.interMessages = [{ role: 'system', content: systemPrompt }];
        this.interMessages.push({ role: 'user', content: `<stage number="1" role="user"><query>${question}</query></stage>` });
        const parser = new XMLParser({
            ignoreAttributes: false,
            attributeNamePrefix: '@_',
            textNodeName: '_text',
            isArray: (name, jpath, isLeafNode, isAttribute) => {
                // Convert tags with the same name to arrays
                return ['query', 'url'].indexOf(name) !== -1;
            },
        });
        const builder = new XMLBuilder({ ignoreAttributes: false, attributeNamePrefix: '@_' });

        let currentAction: string | undefined;

        this.processingInfo = [];

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

            const result = await this.execute(onProcessingUpdate, onAnswerUpdate);
            this.interMessages.push({ role: 'assistant', content: result });

            let parsedResult;
            try {
                parsedResult = parser.parse(result);
            } catch (error) {
                throw new Error(`Error parsing response: ${error}`);
            }

            const stage = parsedResult.stage;

            if (!stage) {
                throw new Error(`Error: No stage found in response`);
            }

            for (const key in stage) {
                if (key === 'thought') {
                    console.log(`Thought: ${stage[key]}`);
                    this.processingNumber++;
                } else if (key === 'action') {
                    currentAction = stage[key] as string;
                    console.log(`Action: ${currentAction}`);
                    if (this.tools[currentAction]) {
                        const nextPrompt = [
                            {
                                type: 'text',
                                text: `<stage number="${i + 1}" role="user">` + builder.build({ action_rules: this.tools[currentAction].getActionRule() }) + `</stage>`,
                            },
                        ];
                        this.interMessages.push({ role: 'user', content: nextPrompt });
                        break;
                    } else {
                        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') {
                    const actionInput = stage[key];
                    console.log(`Action input:`, actionInput);
                    if (currentAction) {
                        try {
                            // Parse the inputs
                            //const parsedInputs = this.parseActionInputs(actionInput.inputs);
                            //console.log(`Parsed inputs:`, parsedInputs);
                            const observation = await this.processAction(currentAction, actionInput.inputs);
                            const nextPrompt = [{ type: 'text', text: `<stage number="${i + 1}" role="user"> <observation>` }, ...observation, { type: 'text', text: '</observation></stage>' }];
                            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') {
                    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 async execute(onProcessingUpdate: (processingUpdate: ProcessingInfo[]) => void, onAnswerUpdate: (answerUpdate: string) => void): Promise<string> {
        const stream = await this.client.chat.completions.create({
            model: 'gpt-4o',
            messages: this.interMessages as ChatCompletionMessageParam[],
            temperature: 0,
            stream: true,
        });

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

        for await (const chunk of stream) {
            let content = chunk.choices[0]?.delta?.content || '';
            fullResponse += content;

            for (const char of content) {
                if (currentTag === 'answer') {
                    currentContent += char;
                    //console.log(char);
                    const streamedAnswer = this.streamedAnswerParser.parse(char);
                    //console.log(streamedAnswer);
                    onAnswerUpdate(streamedAnswer);
                    continue;
                } else if (char === '<') {
                    isInsideTag = true;
                    currentTag = '';
                    currentContent = '';
                } else if (char === '>') {
                    isInsideTag = false;
                    if (currentTag.startsWith('/')) {
                        currentTag = '';
                    }
                } else if (isInsideTag) {
                    currentTag += char;
                } else if (currentTag === 'thought' || currentTag === '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;
    }

    private async processAction(action: string, actionInput: any): Promise<any> {
        if (!(action in this.tools)) {
            throw new Error(`Unknown action: ${action}`);
        }

        const tool = this.tools[action];
        const args: Record<string, any> = {};

        for (const paramName in tool.parameters) {
            if (actionInput[paramName] !== undefined) {
                if (Array.isArray(actionInput[paramName])) {
                    // If the input is already an array, use it as is
                    args[paramName] = actionInput[paramName];
                } else if (typeof actionInput[paramName] === 'object' && actionInput[paramName] !== null) {
                    // If the input is an object, check if it has multiple of the same tag
                    const values = Object.values(actionInput[paramName]);
                    if (values.length > 1) {
                        // If there are multiple values, convert to an array
                        args[paramName] = values;
                    } else {
                        // If there's only one value, use it directly
                        args[paramName] = values[0];
                    }
                } else {
                    // For single values, use them as is
                    args[paramName] = actionInput[paramName];
                }
            } else if (tool.parameters[paramName].required === 'true') {
                throw new Error(`Missing required parameter '${paramName}' for action '${action}'`);
            }
        }

        return await tool.execute(args);
    }

    private parseActionInputs(inputs: any): Record<string, string | string[]> {
        const parsedInputs: Record<string, string | string[]> = {};
        for (const key in inputs) {
            if (Array.isArray(inputs[key])) {
                parsedInputs[key] = inputs[key].map((item: any) => item._text);
            } else {
                parsedInputs[key] = inputs[key]._text;
            }
        }
        return parsedInputs;
    }
}