app/components/MarkdownAnnotator.js

"use client";

import ReactMarkdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
import React from 'react';

// Color mapping for dataset tags
const TAG_COLORS = {
    named: { bg: '#10b98130', border: '#10b981', text: '#34d399', label: 'Named Dataset' },
    descriptive: { bg: '#f59e0b30', border: '#f59e0b', text: '#fbbf24', label: 'Descriptive' },
    vague: { bg: '#a78bfa30', border: '#a78bfa', text: '#c4b5fd', label: 'Vague' },
    'non-dataset': { bg: '#64748b30', border: '#64748b', text: '#94a3b8', label: 'Non-Dataset' },
};

// Strip markdown formatting to get plain text
const stripMd = (s) => s
    .replace(/\*\*\*/g, '')  // bold-italic
    .replace(/\*\*/g, '')     // bold
    .replace(/\*/g, '')       // italic
    .replace(/__/g, '')
    .replace(/_/g, ' ')
    .replace(/^#{1,6}\s+/gm, '')  // headings
    .replace(/\[([^\]]+)\]\([^)]+\)/g, '$1');  // links

// Resolve start/end offsets in rawText based on search term and storedStart
function findNearestOffsets(rawText, searchTerm, storedStart) {
    if (!searchTerm || !rawText) return null;
    
    // Escape regex characters
    const escaped = searchTerm.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
    // Replace whitespace sequences with \s+ to be whitespace-insensitive
    const regexStr = escaped.replace(/\s+/g, '\\s+');
    
    let regex;
    try {
        regex = new RegExp(regexStr, 'gi');
    } catch (e) {
        return null;
    }
    
    const occurrences = [];
    let match;
    while ((match = regex.exec(rawText)) !== null) {
        occurrences.push({
            start: match.index,
            end: match.index + match[0].length,
        });
        if (match[0].length === 0) {
            regex.lastIndex++;
        }
    }
    
    if (occurrences.length === 0) return null;
    
    // Pick the occurrence closest to the stored start position
    let best = occurrences[0];
    let minDiff = Math.abs(best.start - storedStart);
    for (let i = 1; i < occurrences.length; i++) {
        const diff = Math.abs(occurrences[i].start - storedStart);
        if (diff < minDiff) {
            minDiff = diff;
            best = occurrences[i];
        }
    }
    return best;
}

export default function MarkdownAnnotator({ selectedDocIndex, selectedPage, currentPageData, loadingPage, onAnnotate, onTogglePanel, annotationCount, onMentionClick }) {
    const handleMarkdownClick = (e) => {
        const mark = e.target.closest('mark');
        if (mark) {
            const text = mark.textContent || '';
            onMentionClick && onMentionClick(text);
        }
    };
    const handleAnnotateClick = () => {
        const selection = window.getSelection();
        if (selection && selection.toString().trim() !== "" && selection.rangeCount > 0) {
            const text = selection.toString().trim();

            // Compute the character offset of the selection start within the
            // .markdown-preview container. This lets us disambiguate when the
            // same text appears multiple times on the page.
            let selectionOffset = 0;
            const container = document.querySelector('.markdown-preview');
            if (container) {
                try {
                    const range = selection.getRangeAt(0);
                    const preCaretRange = document.createRange();
                    preCaretRange.setStart(container, 0);
                    preCaretRange.setEnd(range.startContainer, range.startOffset);
                    selectionOffset = preCaretRange.toString().length;
                } catch (e) {
                    // Fallback: offset 0 (will just use first occurrence)
                    selectionOffset = 0;
                }
            }

            onAnnotate(text, selectionOffset);
        } else {
            const btn = document.getElementById('annotate-btn');
            if (btn) {
                btn.classList.add('shake');
                setTimeout(() => btn.classList.remove('shake'), 500);
            }
        }
    };

    // Filter out consensus non-datasets (model + judge both agree) and low confidence ones
    const datasets = (currentPageData?.datasets || []).filter(ds => {
        if (ds.dataset_tag === 'non-dataset' && ds.dataset_name?.judge_agrees === true) return false;
        const confidence = ds.dataset_name?.confidence ?? 1.0;
        if (confidence < 0.5) return false;
        return true;
    });
    
    const rawText = currentPageData?.input_text || "";

    // Separate datasets into spanned and text-only
    const spannedDatasets = [];
    const textOnlyDatasets = [];
    datasets.forEach((ds, id) => {
        if (ds.dataset_name?.text && ds.dataset_name.start !== null && ds.dataset_name.end !== null) {
            spannedDatasets.push({ ...ds, _tempId: id });
        } else {
            textOnlyDatasets.push(ds);
        }
    });

    // Translate spans to raw text offsets and generate marker insertions
    const insertions = [];
    spannedDatasets.forEach((ds) => {
        const text = ds.dataset_name.text;
        const start = ds.dataset_name.start;
        const mapped = findNearestOffsets(rawText, text, start);
        if (mapped) {
            const tag = ds.dataset_tag || 'non-dataset';
            insertions.push({ idx: mapped.start, type: 'start', id: ds._tempId, tag });
            insertions.push({ idx: mapped.end, type: 'end', id: ds._tempId, tag });
        } else {
            // Fallback to text-only if mapping fails
            textOnlyDatasets.push(ds);
        }
    });

    // Sort insertions descending.
    // If indices are equal, place 'start' before 'end' in the sorted descending array
    // so that 'start' is processed after 'end' (resulting in START appearing before END in the text)
    insertions.sort((a, b) => {
        if (b.idx !== a.idx) return b.idx - a.idx;
        if (a.type === 'start' && b.type === 'end') return -1;
        if (a.type === 'end' && b.type === 'start') return 1;
        return 0;
    });

    // Inject markers into rawText
    let markedRawText = rawText;
    insertions.forEach(ins => {
        const token = ins.type === 'start' 
            ? `:::MARK_START_${ins.id}_${ins.tag}:::` 
            : `:::MARK_END_${ins.id}_${ins.tag}:::`;
        markedRawText = markedRawText.slice(0, ins.idx) + token + markedRawText.slice(ins.idx);
    });

    // State for tracking active annotations during recursive traversal
    let activeAnnotations = [];

    // Helper to highlight any unspanned text-only datasets via regex match
    const highlightTextOnly = (text) => {
        if (textOnlyDatasets.length === 0 || !text) return text;

        const mentions = textOnlyDatasets
            .filter(ds => ds.dataset_name?.text)
            .map(ds => ({
                name: ds.dataset_name.text,
                tag: ds.dataset_tag || 'non-dataset',
            }))
            .sort((a, b) => b.name.length - a.name.length);

        const seen = new Set();
        const uniqueMentions = mentions.filter(m => {
            if (seen.has(m.name)) return false;
            seen.add(m.name);
            return true;
        });

        if (uniqueMentions.length === 0) return text;

        const escaped = uniqueMentions.map(m => {
            const esc = m.name.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
            return esc.replace(/\s+/g, '\\s+');
        });
        const pattern = new RegExp(`(${escaped.join('|')})`, 'gi');

        const nameToTag = {};
        uniqueMentions.forEach(m => { nameToTag[m.name.toLowerCase()] = m.tag; });

        const parts = text.split(pattern);
        if (parts.length === 1) return text;

        return parts.map((part, i) => {
            const tag = nameToTag[part.toLowerCase()];
            if (tag) {
                const colors = TAG_COLORS[tag] || TAG_COLORS['non-dataset'];
                return (
                    <mark
                        key={i}
                        data-tag={tag}
                        style={{
                            backgroundColor: colors.bg,
                            borderBottom: `2px solid ${colors.border}`,
                            color: colors.text,
                            padding: '1px 3px',
                            borderRadius: '3px',
                            cursor: 'pointer',
                        }}
                        title={`[${colors.label}]`}
                        onClick={() => onMentionClick && onMentionClick(part)}
                    >
                        {part}
                    </mark>
                );
            }
            return part;
        });
    };

    const markerRegex = /(:::MARK_START_\d+_[a-zA-Z0-9-]+:::|:::MARK_END_\d+_[a-zA-Z0-9-]+:::)/g;

    // Parse marker tokens and split text leaf nodes
    const splitTextByMarkers = (text) => {
        if (!text) return text;
        const parts = text.split(markerRegex);
        if (parts.length === 1) {
            if (activeAnnotations.length > 0) {
                const active = activeAnnotations[activeAnnotations.length - 1];
                const colors = TAG_COLORS[active.tag] || TAG_COLORS['non-dataset'];
                return (
                    <mark
                        data-tag={active.tag}
                        style={{
                            backgroundColor: colors.bg,
                            borderBottom: `2px solid ${colors.border}`,
                            color: colors.text,
                            padding: '1px 3px',
                            borderRadius: '3px',
                            cursor: 'pointer',
                        }}
                        title={`[${colors.label}]`}
                        onClick={() => onMentionClick && onMentionClick(text)}
                    >
                        {text}
                    </mark>
                );
            }
            return highlightTextOnly(text);
        }

        return parts.map((part, i) => {
            if (part.startsWith(':::MARK_START_')) {
                const match = part.match(/:::MARK_START_(\d+)_([a-zA-Z0-9-]+):::/);
                if (match) {
                    const id = parseInt(match[1], 10);
                    const tag = match[2];
                    activeAnnotations.push({ id, tag });
                }
                return null;
            }
            if (part.startsWith(':::MARK_END_')) {
                const match = part.match(/:::MARK_END_(\d+)_([a-zA-Z0-9-]+):::/);
                if (match) {
                    const id = parseInt(match[1], 10);
                    activeAnnotations = activeAnnotations.filter(ann => ann.id !== id);
                }
                return null;
            }
            
            if (part === '') return null;
            
            if (activeAnnotations.length > 0) {
                const active = activeAnnotations[activeAnnotations.length - 1];
                const colors = TAG_COLORS[active.tag] || TAG_COLORS['non-dataset'];
                return (
                    <mark
                        key={i}
                        data-tag={active.tag}
                        style={{
                            backgroundColor: colors.bg,
                            borderBottom: `2px solid ${colors.border}`,
                            color: colors.text,
                            padding: '1px 3px',
                            borderRadius: '3px',
                            cursor: 'pointer',
                        }}
                        title={`[${colors.label}]`}
                        onClick={() => onMentionClick && onMentionClick(part)}
                    >
                        {part}
                    </mark>
                );
            }
            
            return highlightTextOnly(part);
        });
    };

    // Recursive helper: processes children at any depth so text inside
    // <strong>, <em>, <a>, etc. also gets highlighted.
    const processChildren = (children) =>
        React.Children.map(children, child => {
            if (typeof child === 'string') {
                return splitTextByMarkers(child);
            }
            if (React.isValidElement(child) && child.props?.children) {
                return React.cloneElement(child, {}, processChildren(child.props.children));
            }
            return child;
        });

    // Build component overrides for all block-level and inline elements
    const highlightWrapper = (Tag) => ({ children, ...props }) => (
        <Tag {...props}>{processChildren(children)}</Tag>
    );

    const highlightComponents = {
        p: highlightWrapper('p'),
        li: highlightWrapper('li'),
        td: highlightWrapper('td'),
        th: highlightWrapper('th'),
        h1: highlightWrapper('h1'),
        h2: highlightWrapper('h2'),
        h3: highlightWrapper('h3'),
        h4: highlightWrapper('h4'),
        h5: highlightWrapper('h5'),
        h6: highlightWrapper('h6'),
        blockquote: highlightWrapper('blockquote'),
        strong: highlightWrapper('strong'),
        em: highlightWrapper('em'),
    };

    return (
        <div className="annotator-container">
            <div className="annotator-header">
                <h2>Markdown Annotation</h2>
                <button
                    id="annotate-btn"
                    onClick={handleAnnotateClick}
                    className="btn btn-primary"
                    title="Select text below, then click to annotate"
                >
                    ✍️ Annotate Selection
                </button>
            </div>

            {/* Dataset legend */}
            {datasets.length > 0 && (
                <div className="dataset-legend">
                    {Object.entries(TAG_COLORS).map(([tag, colors]) => {
                        const count = datasets.filter(ds => ds.dataset_tag === tag).length;
                        if (count === 0) return null;
                        return (
                            <span key={tag} className="legend-item" style={{ borderColor: colors.border }}>
                                <span className="legend-dot" style={{ backgroundColor: colors.border }} />
                                {colors.label} ({count})
                            </span>
                        );
                    })}
                </div>
            )}

            <div className="markdown-content">
                <div className="markdown-content-header">
                    <h3>Page {selectedPage !== null ? selectedPage + 1 : ''}</h3>
                    {datasets.length > 0 && (
                        <span className="dataset-count">
                            {datasets.length} data mention{datasets.length !== 1 ? 's' : ''} detected
                        </span>
                    )}
                </div>

                {loadingPage ? (
                    <div className="loading-spinner-container">
                        <div className="loading-spinner" />
                        <p>Loading page data...</p>
                    </div>
                ) : currentPageData ? (
                    <div className="markdown-preview" onClick={handleMarkdownClick}>
                        <ReactMarkdown
                            remarkPlugins={[remarkGfm]}
                            components={highlightComponents}
                        >
                            {markedRawText || "No text available."}
                        </ReactMarkdown>
                    </div>
                ) : (
                    <p className="text-muted">Select a document and page to view extracted text.</p>
                )}
            </div>
        </div>
    );
}