import { action, computed, makeObservable, observable, observe, reaction, runInAction, ObservableSet } from 'mobx';
import { observer } from 'mobx-react';
import OpenAI, { ClientOptions } from 'openai';
import * as React from 'react';
import { Doc, DocListCast } from '../../../../fields/Doc';
import { CsvCast, DocCast, PDFCast, StrCast } from '../../../../fields/Types';
import { DocumentType } from '../../../documents/DocumentTypes';
import { Docs } from '../../../documents/Documents';
import { LinkManager } from '../../../util/LinkManager';
import { ViewBoxAnnotatableComponent } from '../../DocComponent';
import { FieldView, FieldViewProps } from '../FieldView';
import './ChatBox.scss';
import MessageComponentBox from './MessageComponent';
import { ASSISTANT_ROLE, AssistantMessage, AI_Document, Citation, CHUNK_TYPE, RAGChunk, getChunkType, TEXT_TYPE, SimplifiedChunk } from './types';
import { Vectorstore } from './vectorstore/Vectorstore';
import { Agent } from './Agent';
import dotenv from 'dotenv';
import { DocData, DocViews } from '../../../../fields/DocSymbols';
import { AnswerParser } from './AnswerParser';
import { DocumentManager } from '../../../util/DocumentManager';
import { v4 as uuidv4 } from 'uuid';
import { chunk } from 'lodash';
import { DocUtils } from '../../../documents/DocUtils';

dotenv.config();

@observer
export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
    @observable history: AssistantMessage[] = [];
    @observable.deep current_message: AssistantMessage | undefined = undefined;

    @observable isLoading: boolean = false;
    @observable isUploadingDocs: boolean = false;
    @observable expandedScratchpadIndex: number | null = null;
    @observable inputValue: string = '';
    @observable private linked_docs_to_add: ObservableSet<Doc> = observable.set();
    @observable private linked_csv_files: { filename: string; id: string; text: string }[] = [];
    private openai: OpenAI;
    private vectorstore_id: string;
    private documents: AI_Document[] = [];
    private _oldWheel: any;
    private vectorstore: Vectorstore;
    private agent: Agent; // Add the ChatBot instance

    public static LayoutString(fieldKey: string) {
        return FieldView.LayoutString(ChatBox, fieldKey);
    }

    constructor(props: FieldViewProps) {
        super(props);
        makeObservable(this);
        this.openai = this.initializeOpenAI();
        if (StrCast(this.dataDoc.vectorstore_id) == '') {
            console.log('new_id');
            this.vectorstore_id = uuidv4();
            this.dataDoc.vectorstore_id = this.vectorstore_id;
        } else {
            this.vectorstore_id = StrCast(this.dataDoc.vectorstore_id);
        }
        this.vectorstore = new Vectorstore(this.vectorstore_id, this.retrieveDocIds);
        this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory, this.retrieveCSVData, this.addLinkedUrlDoc);

        reaction(
            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, content: msg.content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
            serializableHistory => {
                this.dataDoc.data = JSON.stringify(serializableHistory);
            }
        );
    }

    @action
    addDocToVectorstore = async (newLinkedDoc: Doc) => {
        await this.vectorstore.addAIDoc(newLinkedDoc);
    };

    @action
    addCSVForAnalysis = async (newLinkedDoc: Doc) => {
        console.log('adding csv file for analysis');
        if (!newLinkedDoc.chunk_simpl) {
            const csvData: string = StrCast(newLinkedDoc.text);
            console.log('CSV Data:', csvData);
            const completion = await this.openai.chat.completions.create({
                messages: [
                    {
                        role: 'system',
                        content:
                            'You are an AI assistant tasked with summarizing the content of a CSV file. You will be provided with the data from the CSV file and your goal is to generate a concise summary that captures the main themes, trends, and key points represented in the data.',
                    },
                    {
                        role: 'user',
                        content: `Please provide a comprehensive summary of the CSV file based on the provided data. Ensure the summary highlights the most important information, patterns, and insights. Your response should be in paragraph form and be concise.
                
                    CSV Data:
                
                    ${csvData}
                
                    **********
                    Summary:`,
                    },
                ],
                model: 'gpt-3.5-turbo',
            });
            console.log('CSV Data:', csvData);
            const csvId = uuidv4();

            this.linked_csv_files.push({
                filename: CsvCast(newLinkedDoc.data).url.pathname,
                id: csvId,
                text: csvData,
            });

            console.log(this.linked_csv_files);
            const chunkToAdd = {
                chunkId: csvId,
                chunkType: CHUNK_TYPE.CSV,
            };
            newLinkedDoc.chunk_simpl = JSON.stringify({ chunks: [chunkToAdd] });
            newLinkedDoc.summary = completion.choices[0].message.content!;
        }
    };

    @action
    toggleToolLogs = (index: number) => {
        this.expandedScratchpadIndex = this.expandedScratchpadIndex === index ? null : index;
    };

    initializeOpenAI() {
        console.log(process.env.OPENAI_KEY);
        const configuration: ClientOptions = {
            apiKey: process.env.OPENAI_KEY,
            dangerouslyAllowBrowser: true,
        };
        return new OpenAI(configuration);
    }

    onPassiveWheel = (e: WheelEvent) => {
        if (this._props.isContentActive()) {
            e.stopPropagation();
        }
    };

    @action
    askGPT = async (event: React.FormEvent<HTMLFormElement>): Promise<void> => {
        event.preventDefault();
        this.inputValue = '';

        const textInput = event.currentTarget.elements.namedItem('messageInput') as HTMLInputElement;
        const trimmedText = textInput.value.trim();

        if (trimmedText) {
            try {
                textInput.value = '';
                this.history.push({ role: ASSISTANT_ROLE.USER, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: trimmedText, citation_ids: null }] });
                this.isLoading = true;
                this.current_message = { role: ASSISTANT_ROLE.ASSISTANT, content: [], thoughts: [], actions: [], citations: [] };

                const onUpdate = (update: AssistantMessage) => {
                    runInAction(() => {
                        this.current_message = { ...update };
                    });
                };

                const finalMessage = await this.agent.askAgent(trimmedText, 20, onUpdate);

                runInAction(() => {
                    this.history.push({ ...finalMessage });
                    this.current_message = undefined;
                    this.dataDoc.data = JSON.stringify(this.history);
                });
            } catch (err) {
                console.error('Error:', err);
                this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }] });
            } finally {
                this.isLoading = false;
            }
        }
    };

    @action
    updateMessageCitations = (index: number, citations: Citation[]) => {
        if (this.history[index]) {
            this.history[index].citations = citations;
        }
    };

    @action
    addLinkedUrlDoc = async (url: string, id: string) => {
        const doc = Docs.Create.WebDocument(url);

        const linkDoc = Docs.Create.LinkDocument(this.Document, doc);
        LinkManager.Instance.addLink(linkDoc);

        const chunkToAdd = {
            chunkId: id,
            chunkType: CHUNK_TYPE.URL,
        };

        doc.chunk_simpl = JSON.stringify({ chunks: [chunkToAdd] });
    };

    @action
    handleCitationClick = (citation: Citation) => {
        console.log('Citation clicked:', citation);
        const currentLinkedDocs: Doc[] = this.linkedDocs;

        const chunkId = citation.chunk_id;

        for (let doc of currentLinkedDocs) {
            if (doc.chunk_simpl) {
                const docChunkSimpl = JSON.parse(StrCast(doc.chunk_simpl)) as { chunks: SimplifiedChunk[] };
                console.log(docChunkSimpl);
                const foundChunk = docChunkSimpl.chunks.find(chunk => chunk.chunkId === chunkId);
                if (foundChunk) {
                    switch (getChunkType(foundChunk.chunkType)) {
                        case CHUNK_TYPE.IMAGE:
                            const values = foundChunk.location?.replace(/[\[\]]/g, '').split(',');

                            if (values?.length !== 4) {
                                console.error('Location string must contain exactly 4 numbers');
                                return;
                            }

                            const x1 = parseFloat(values[0]) * Doc.NativeWidth(doc);
                            const y1 = parseFloat(values[1]) * Doc.NativeHeight(doc);
                            const x2 = parseFloat(values[2]) * Doc.NativeWidth(doc);
                            const y2 = parseFloat(values[3]) * Doc.NativeHeight(doc);

                            const annotationKey = Doc.LayoutFieldKey(doc) + '_annotations';

                            const existingDoc = DocListCast(doc[DocData][annotationKey]).find(d => d.citation_id === citation.citation_id);
                            const highlightDoc = existingDoc ?? this.createImageCitationHighlight(x1, y1, x2, y2, citation, annotationKey, doc);

                            DocumentManager.Instance.showDocument(highlightDoc, { willZoomCentered: true }, () => {});
                            break;
                        case CHUNK_TYPE.TEXT:
                            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                                const firstView = Array.from(doc[DocViews])[0];
                                firstView.ComponentView?.search?.(citation.direct_text);
                            });
                            break;
                        case CHUNK_TYPE.URL:
                            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                                const firstView = Array.from(doc[DocViews])[0];
                            });
                            break;
                        case CHUNK_TYPE.CSV:
                            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                                const firstView = Array.from(doc[DocViews])[0];
                            });
                            break;
                        default:
                            console.log('Chunk type not supported');
                            break;
                    }
                }
            }
        }
    };

    createImageCitationHighlight = (x1: number, y1: number, x2: number, y2: number, citation: Citation, annotationKey: string, pdfDoc: Doc): Doc => {
        const highlight_doc = Docs.Create.FreeformDocument([], {
            x: x1,
            y: y1,
            _width: x2 - x1,
            _height: y2 - y1,
            backgroundColor: 'rgba(255, 255, 0, 0.5)',
        });
        highlight_doc[DocData].citation_id = citation.citation_id;
        Doc.AddDocToList(pdfDoc[DocData], annotationKey, highlight_doc);
        highlight_doc.annotationOn = pdfDoc;
        Doc.SetContainer(highlight_doc, pdfDoc);
        return highlight_doc;
    };

    componentDidMount() {
        this._props.setContentViewBox?.(this);
        if (this.dataDoc.data) {
            try {
                const storedHistory = JSON.parse(StrCast(this.dataDoc.data));
                runInAction(() => {
                    this.history.push(
                        ...storedHistory.map((msg: AssistantMessage) => ({
                            role: msg.role,
                            content: msg.content,
                            follow_up_questions: msg.follow_up_questions,
                            citations: msg.citations,
                        }))
                    );
                });
            } catch (e) {
                console.error('Failed to parse history from dataDoc:', e);
            }
        } else {
            runInAction(() => {
                this.history.push({
                    role: ASSISTANT_ROLE.ASSISTANT,
                    content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.', citation_ids: null }],
                });
            });
        }
        reaction(
            () => {
                const linkedDocs = LinkManager.Instance.getAllRelatedLinks(this.Document)
                    .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
                    .map(d => DocCast(d?.annotationOn, d))
                    .filter(d => d);
                return linkedDocs;
            },

            linked => linked.forEach(doc => this.linked_docs_to_add.add(doc))
        );

        observe(this.linked_docs_to_add, change => {
            if (change.type === 'add') {
                runInAction(() => {
                    this.isUploadingDocs = true;
                });
                if (PDFCast(change.newValue.data)) {
                    this.addDocToVectorstore(change.newValue);
                } else if (CsvCast(change.newValue.data)) {
                    this.addCSVForAnalysis(change.newValue);
                }
                runInAction(() => {
                    this.isUploadingDocs = false;
                });
            } else if (change.type === 'delete') {
                console.log('Deleted docs: ', change.oldValue);
            }
        });
    }

    @computed
    get linkedDocs() {
        return LinkManager.Instance.getAllRelatedLinks(this.Document)
            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
            .map(d => DocCast(d?.annotationOn, d))
            .filter(d => d);
    }

    @computed
    get docIds() {
        return LinkManager.Instance.getAllRelatedLinks(this.Document)
            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
            .map(d => DocCast(d?.annotationOn, d))
            .filter(d => d)
            .filter(d => d.ai_doc_id)
            .map(d => StrCast(d.ai_doc_id));
    }

    @computed
    get summaries(): string {
        return (
            LinkManager.Instance.getAllRelatedLinks(this.Document)
                .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
                .map(d => DocCast(d?.annotationOn, d))
                .filter(d => d)
                .filter(d => d.summary)
                .map((doc, index) => {
                    if (PDFCast(doc.data)) {
                        return `<summary file_name="${PDFCast(doc.data).url.pathname}" applicable_tools=["rag"]>${doc.summary}</summary>`;
                    } else if (CsvCast(doc.data)) {
                        return `<summary file_name="${CsvCast(doc.data).url.pathname}" applicable_tools=["dataAnalysis"]>${doc.summary}</summary>`;
                    } else {
                        return `${index + 1}) ${doc.summary}`;
                    }
                })
                .join('\n') + '\n'
        );
    }

    @computed
    get linkedCSVs(): { filename: string; id: string; text: string }[] {
        return this.linked_csv_files;
    }

    @computed
    get formattedHistory(): string {
        let history = '<chat_history>\n';
        for (const message of this.history) {
            history += `<${message.role}>${message.content.map(content => content.text).join(' ')}</${message.role}>\n`;
        }
        history += '</chat_history>';
        return history;
    }

    retrieveSummaries = () => {
        return this.summaries;
    };

    retrieveCSVData = () => {
        return this.linkedCSVs;
    };

    retrieveFormattedHistory = () => {
        return this.formattedHistory;
    };

    retrieveDocIds = () => {
        return this.docIds;
    };

    @action
    handleFollowUpClick = (question: string) => {
        console.log('Follow-up question clicked:', question);
        this.inputValue = question;
    };
    render() {
        return (
            <div className="chatBox">
                {this.isUploadingDocs && <div className="uploading-overlay"></div>}
                <div
                    className="scroll-box chat-content"
                    ref={r => {
                        this._oldWheel?.removeEventListener('wheel', this.onPassiveWheel);
                        this._oldWheel = r;
                        r?.addEventListener('wheel', this.onPassiveWheel, { passive: false });
                    }}>
                    <div className="messages">
                        {this.history.map((message, index) => (
                            <MessageComponentBox key={index} message={message} index={index} onFollowUpClick={this.handleFollowUpClick} onCitationClick={this.handleCitationClick} updateMessageCitations={this.updateMessageCitations} />
                        ))}
                        {!this.current_message ? null : (
                            <MessageComponentBox
                                key={this.history.length}
                                message={this.current_message}
                                index={this.history.length}
                                onFollowUpClick={this.handleFollowUpClick}
                                onCitationClick={this.handleCitationClick}
                                updateMessageCitations={this.updateMessageCitations}
                            />
                        )}
                    </div>
                </div>
                <form onSubmit={this.askGPT} className="chat-form">
                    <input type="text" name="messageInput" autoComplete="off" placeholder="Type a message..." value={this.inputValue} onChange={e => (this.inputValue = e.target.value)} />
                    <button type="submit" disabled={this.isLoading}>
                        {this.isLoading ? 'Thinking...' : 'Send'}
                    </button>
                </form>
            </div>
        );
    }
}

Docs.Prototypes.TemplateMap.set(DocumentType.CHAT, {
    layout: { view: ChatBox, dataField: 'data' },
    options: { acl: '', chat: '', chat_history: '', chat_thread_id: '', chat_assistant_id: '', chat_vector_store_id: '' },
});
