import OpenAI from 'openai';
import { Tool, AgentMessage, AssistantMessage, TEXT_TYPE, CHUNK_TYPE, ASSISTANT_ROLE } 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 { StreamParser } from './StreamParser';
import { v4 as uuidv4 } from 'uuid';
import { AnswerParser } from './AnswerParser';

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;

    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, maxTurns: number = 30, onUpdate: (update: AssistantMessage) => void): 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: '@_' });
        const builder = new XMLBuilder({ ignoreAttributes: false, attributeNamePrefix: '@_' });
        let currentAction: string | undefined;

        let assistantMessage: AssistantMessage = {
            role: ASSISTANT_ROLE.ASSISTANT,
            content: [],
            thoughts: [],
            actions: [],
            citations: [],
        };

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

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

            let parsedResult;
            try {
                parsedResult = parser.parse(result);
            } catch (error) {
                console.log('Error: Invalid XML response from bot');
                assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
                return assistantMessage;
            }

            const stage = parsedResult.stage;

            if (!stage) {
                console.log('Error: No stage found in response');
                assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
                return assistantMessage;
            }

            for (const key in stage) {
                if (!assistantMessage.actions) {
                    assistantMessage.actions = [];
                }
                if (key === 'thought') {
                    console.log(`Thought: ${stage[key]}`);
                    this.thoughtNumber++;
                } else if (key === 'action') {
                    currentAction = stage[key] as string;
                    console.log(`Action: ${currentAction}`);
                    onUpdate({ ...assistantMessage });
                    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 = builder.build({ action_input: stage[key] });
                    console.log(`Action input: ${actionInput}`);
                    if (currentAction) {
                        try {
                            const observation = await this.processAction(currentAction, stage[key]);
                            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.actionNumber++; //might not work with no tool
                            break;
                        } catch (error) {
                            console.log(`Error processing action: ${error}`);
                            assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
                            return assistantMessage;
                        }
                    } else {
                        console.log('Error: Action input without a valid action');
                        assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
                        return assistantMessage;
                    }
                } else if (key === 'answer') {
                    console.log('Answer found. Ending query.');
                    const parsedAnswer = AnswerParser.parse(result, assistantMessage);
                    onUpdate({ ...parsedAnswer });
                    return parsedAnswer;
                }
            }
        }
        console.log('Reached maximum turns. Ending query.');
        return assistantMessage;
    }

    private async execute(assistantMessage: AssistantMessage, onUpdate: (update: AssistantMessage) => 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;
        let isInsideActionInput: boolean = false;
        let actionInputContent: string = '';

        if (!assistantMessage.actions) {
            assistantMessage.actions = [];
        }

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

            for (const char of content) {
                if (char === '<') {
                    isInsideTag = true;
                    if (currentTag && currentContent) {
                        if (currentTag === 'action_input') {
                            assistantMessage.actions[assistantMessage.actions.length - 1].action_input = actionInputContent;
                            actionInputContent = '';
                        } else {
                            this.processStreamedContent(currentTag, currentContent, assistantMessage);
                        }
                        onUpdate({ ...assistantMessage });
                    }
                    currentTag = '';
                    currentContent = '';
                } else if (char === '>') {
                    isInsideTag = false;
                    if (currentTag === 'action_input') {
                        isInsideActionInput = true;
                    } else if (currentTag === '/action_input') {
                        isInsideActionInput = false;
                        console.log('Action input:', actionInputContent);
                        assistantMessage.actions[assistantMessage.actions.length - 1].action_input = actionInputContent;
                        actionInputContent = '';
                        onUpdate({ ...assistantMessage });
                    }
                    if (currentTag.startsWith('/')) {
                        currentTag = '';
                    }
                } else if (isInsideTag) {
                    currentTag += char;
                } else if (isInsideActionInput) {
                    actionInputContent += char;
                } else {
                    currentContent += char;
                    if (currentTag === 'thought' || currentTag === 'action') {
                        this.processStreamedContent(currentTag, currentContent, assistantMessage);
                        onUpdate({ ...assistantMessage });
                    }
                }
            }
        }

        return fullResponse;
    }

    private processStreamedContent(tag: string, content: string, assistantMessage: AssistantMessage) {
        if (!assistantMessage.thoughts) {
            assistantMessage.thoughts = [];
        }
        if (!assistantMessage.actions) {
            assistantMessage.actions = [];
        }
        switch (tag) {
            case 'thought':
                assistantMessage.thoughts[this.thoughtNumber] = content;
                break;
            case 'action':
                assistantMessage.actions[this.actionNumber].action = content;

                break;
            case 'action_input':
                assistantMessage.actions[this.actionNumber].action_input = content;
                break;
        }
    }

    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) {
                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);
    }
}