public/elements/Message.jsx

import React, { useState } from 'react';
import { Button } from "@/components/ui/button"
import {
    DropdownMenu,
    DropdownMenuContent,
    DropdownMenuGroup,
    DropdownMenuItem,
    DropdownMenuShortcut,
    DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu"

import { Eye, EyeClosed } from "lucide-react"
import {
    Tooltip,
    TooltipContent,
    TooltipTrigger,
} from "@/components/ui/tooltip"


import { Button } from "@/components/ui/button"
// Internal Badge component to avoid import errors
const Badge = ({ children, className }) => (
    <span className={` py-0.5 transition-colors text-main-700 ${className}`}>
        {children}
    </span>
);

function MessageWithHighlight({ message, dic, show }) {
    if (!message) return null;

    const [localMessage, setLocalMessage] = useState(message);
    const [open, setOpen] = useState(false);
    const [position, setPosition] = useState({ x: 0, y: 0 });
    const [selectionData, setSelectionData] = useState(null);

    const handleMouseUp = () => {
        const selection = window.getSelection();

        // 1. Check basic validity
        if (!selection || selection.rangeCount === 0 || selection.isCollapsed) {
            // If just a click (collapsed), we assume menu shouldn't open or should close
            // But we don't want to close immediately if clicking inside the menu, handled by DropdownMenu logic
            return;
        }

        // 2. Expand Selection to Word Boundaries
        const anchorNode = selection.anchorNode;
        const focusNode = selection.focusNode;

        // Only support single text node selection for simplicity to map back to data-index
        if (anchorNode !== focusNode || anchorNode.nodeType !== 3) {
            setOpen(false);
            return;
        }

        const textContent = anchorNode.textContent;
        let start = Math.min(selection.anchorOffset, selection.focusOffset);
        let end = Math.max(selection.anchorOffset, selection.focusOffset);

        // Expand left
        while (start > 0 && /\S/.test(textContent[start - 1])) {
            start--;
        }
        // Expand right
        while (end < textContent.length && /\S/.test(textContent[end])) {
            end++;
        }

        // Update visual selection
        const newRange = document.createRange();
        newRange.setStart(anchorNode, start);
        newRange.setEnd(anchorNode, end);
        selection.removeAllRanges();
        selection.addRange(newRange);

        const newText = textContent.slice(start, end);
        if (!newText.trim()) return;

        // 3. Validation Logic
        // Check 1: Less than 4 words
        const words = newText.split(/\s+/);
        if (words.length >= 4) {
            setOpen(false);
            return;
        }

        // Check 2: Does not contain highlighted word (brackets)
        if (newText.includes('[') || newText.includes(']')) {
            setOpen(false);
            return;
        }

        // Check 3: Check if selection is inside a Badge
        const parentElement = anchorNode.parentElement;
        if (parentElement.closest('.bg-pii')) {
            setOpen(false);
            return;
        }

        // Identify Part Index
        const partIndex = parseInt(parentElement.getAttribute('data-index'), 10);
        if (isNaN(partIndex)) {
            setOpen(false);
            return;
        }

        const rect = newRange.getBoundingClientRect();
        setPosition({
            x: rect.left,
            y: rect.bottom
        });
        setSelectionData({
            partIndex,
            start,
            end,
            text: newText
        });
        setOpen(true);
    };

    const handleAnonymize = () => {
        if (!selectionData) return;
        const { partIndex, start, end, text } = selectionData;

        // Split by brackets to reconstruct parts array as seen in render
        const parts = localMessage.split(/(\[[^\]]+\])/g);

        if (parts[partIndex] === undefined) return;

        const originalPart = parts[partIndex];
        const newPart = originalPart.substring(0, start) + `[ANONYME_${text}]` + originalPart.substring(end);

        parts[partIndex] = newPart;
        const newMessage = parts.join('');

        setLocalMessage(newMessage);
        setOpen(false);
        // Clear selection
        window.getSelection().removeAllRanges();
    };

    // Split by brackets, keeping the brackets in the split result for now
    // Regex: split by [content]
    const parts = localMessage.split(/(\[[^\]]+\])/g);

    return (
        <div onMouseUp={handleMouseUp} className="relative px-5 py-2.5">
            <p className="leading-[1.75] whitespace-pre-wrap">
                {parts.map((part, index) => {
                    // Check if part is [something]
                    const match = part.match(/^\[(\w+)_(.*?)\]$/);

                    if (match && dic[match[2]]) {
                        return (
                            show ? <Badge key={index} className="bg-pii text-blue-800 rounded-full px-2 py-1 mx-1 cursor-pointer">
                                {dic[match[2]].value}
                            </Badge> :
                                <Tooltip>
                                    <TooltipTrigger asChild>
                                        <button className=""><Badge key={index} className="bg-pii text-blue-800 rounded-sm px-1 py-1 mx-1 cursor-pointer">
                                            {part.toUpperCase()}
                                        </Badge></button>
                                    </TooltipTrigger>
                                    <TooltipContent>
                                        <p>{dic[match[2]].value}</p>
                                    </TooltipContent>
                                </Tooltip>
                        );
                    }
                    return <span key={index} data-index={index}>{part}</span>;
                })}
            </p>

            <DropdownMenu open={open} onOpenChange={setOpen}>
                <DropdownMenuTrigger asChild>
                    <div
                        style={{
                            position: 'fixed',
                            top: position.y,
                            left: position.x,
                            width: '1px',
                            height: '1px',
                            visibility: 'hidden',
                            pointerEvents: 'none'
                        }}
                    />
                </DropdownMenuTrigger>
                <DropdownMenuContent align="start">
                    <DropdownMenuGroup>
                        <DropdownMenuItem onSelect={handleAnonymize}>
                            Anonymiser
                        </DropdownMenuItem>
                    </DropdownMenuGroup>
                </DropdownMenuContent>
            </DropdownMenu>
        </div>
    );
}

export default function Message() {
    const [show, setShow] = useState(true);
    return (
        <div className="flex flex-col gap-2">
            {/* <ReactMarkdown> */}
            <MessageWithHighlight message={props.message} dic={props.dic} show={show} />
            {/* </ReactMarkdown> */}

            {/* <div className='flex flex-row-reverse w-full'>
                <div className="relative bg-accent rounded-3xl px-5 py-2.5 user-message flex-grow-0">
                    <Markdown dic={props.dic} show={show}>{props.message}</Markdown>
                </div>
            </div> */}
            <div className="flex justify-end">
                <Button variant="outline" size="icon" onClick={() => setShow(!show)}>
                    {show ? <Eye /> : <EyeClosed />}
                </Button>
            </div>
        </div>
    );
}