import { ASSISTANT_ROLE, AssistantMessage, Citation, CHUNK_TYPE, TEXT_TYPE, 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 index="([^"]+)" chunk_id="([^"]+)" type="([^"]+)">([\s\S]*?)<\/citation>/g;
        const followUpQuestionsRegex = /<follow_up_questions>([\s\S]*?)<\/follow_up_questions>/;
        const questionRegex = /<question>(.*?)<\/question>/g;
        const groundedTextRegex = /<grounded_text citation_index="([^"]+)">([\s\S]*?)<\/grounded_text>/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();
        let content: AssistantMessage['content'] = [];
        let citations: Citation[] = [];
        let contentIndex = 0;

        // Parse citations
        let citationMatch;
        const citationMap = new Map<string, string>();
        while ((citationMatch = citationRegex.exec(rawTextContent)) !== null) {
            const [_, index, chunk_id, type, direct_text] = citationMatch;
            const citation_id = uuid();
            citationMap.set(index, citation_id);
            citations.push({
                direct_text: direct_text.trim(),
                type: getChunkType(type),
                chunk_id,
                citation_id,
            });
        }

        // Parse grounded text content
        const parseGroundedText = (text: string): AssistantMessage['content'] => {
            const result: AssistantMessage['content'] = [];
            let lastIndex = 0;
            let match;

            while ((match = groundedTextRegex.exec(text)) !== null) {
                const [fullMatch, citationIndex, groundedText] = match;
                const citation_ids = citationIndex.split(',').map(index => citationMap.get(index) || '');

                result.push({
                    index: contentIndex++,
                    type: TEXT_TYPE.GROUNDED,
                    text: groundedText.trim(),
                    citation_ids,
                });

                lastIndex = match.index + fullMatch.length;
            }

            return result;
        };

        content = parseGroundedText(rawTextContent);

        let followUpQuestions: string[] = [];
        if (followUpQuestionsMatch) {
            const questionsText = followUpQuestionsMatch[1];
            let questionMatch;
            while ((questionMatch = questionRegex.exec(questionsText)) !== null) {
                followUpQuestions.push(questionMatch[1].trim());
            }
        }

        const assistantResponse: AssistantMessage = {
            role: ASSISTANT_ROLE.ASSISTANT,
            content,
            follow_up_questions: followUpQuestions,
            citations,
        };

        return assistantResponse;
    }
}