import { ASSISTANT_ROLE, AssistantMessage, Citation, getChunkType } from './types';
import { v4 as uuid } from 'uuid';

export class AnswerParser {
    static parse(xml: string): AssistantMessage {
        const answerRegex = /<answer>([\s\S]*?)<\/answer>/;
        const citationRegex = /<citation chunk_id="([^"]+)" type="([^"]+)">(.*?)<\/citation>/g;
        const followUpQuestionsRegex = /<follow_up_questions>([\s\S]*?)<\/follow_up_questions>/;
        const questionRegex = /<question>(.*?)<\/question>/g;

        const answerMatch = answerRegex.exec(xml);
        const followUpQuestionsMatch = followUpQuestionsRegex.exec(xml);

        if (!answerMatch) {
            throw new Error('Invalid XML: Missing <answer> tag.');
        }

        const rawTextContent = answerMatch[1].trim();
        const textContentWithCitations = rawTextContent.replace(citationRegex, '');
        const textContent = textContentWithCitations.replace(followUpQuestionsRegex, '').trim();

        let citations: Citation[] = [];
        let match: RegExpExecArray | null;

        let plainTextOffset = 0;
        let citationOffset = 0;

        while ((match = citationRegex.exec(rawTextContent)) !== null) {
            const [fullMatch, chunk_id, type, direct_text] = match;
            const citationStartIndex = match.index;
            const citationPlainStart = citationStartIndex - citationOffset;

            citations.push({
                direct_text: direct_text.trim(),
                type: getChunkType(type),
                chunk_id: chunk_id,
                text_location: citationPlainStart,
                citation_id: uuid(),
            });

            citationOffset += fullMatch.length;
        }

        let followUpQuestions: string[] = [];
        if (followUpQuestionsMatch) {
            const questionsText = followUpQuestionsMatch[1];
            let questionMatch: RegExpExecArray | null;

            while ((questionMatch = questionRegex.exec(questionsText)) !== null) {
                followUpQuestions.push(questionMatch[1].trim());
            }
        }
        const assistantResponse: AssistantMessage = {
            role: ASSISTANT_ROLE.ASSISTANT,
            text_content: textContent,
            follow_up_questions: followUpQuestions,
            citations: citations,
        };

        return assistantResponse;
    }
}