stable-diffusion.cpp-rest/webui/components/prompt-textarea.tsx

'use client';

import React, { useState, useRef, useEffect } from 'react';
import { cn } from '@/lib/utils';

interface PromptTextareaProps {
  value: string;
  onChange: (value: string) => void;
  placeholder?: string;
  className?: string;
  rows?: number;
  loras?: string[];
  embeddings?: string[];
}

interface Suggestion {
  text: string;
  type: 'lora' | 'embedding';
  displayText: string;
}

export function PromptTextarea({
  value,
  onChange,
  placeholder,
  className,
  rows = 3,
  loras = [],
  embeddings = [],
}: PromptTextareaProps) {
  const textareaRef = useRef<HTMLTextAreaElement>(null);
  const [highlighted, setHighlighted] = useState<React.ReactNode[]>([]);
  const [suggestions, setSuggestions] = useState<Suggestion[]>([]);
  const [showSuggestions, setShowSuggestions] = useState(false);
  const [selectedIndex, setSelectedIndex] = useState(0);
  const [cursorPosition, setCursorPosition] = useState(0);
  const suggestionsRef = useRef<HTMLDivElement>(null);
  const justInsertedRef = useRef(false);

  useEffect(() => {
    highlightSyntax(value);

    // Skip suggestions if we just inserted one
    if (justInsertedRef.current) {
      justInsertedRef.current = false;
      return;
    }

    updateSuggestions(value, cursorPosition);
  }, [value, loras, embeddings, cursorPosition]);

  const updateSuggestions = (text: string, position: number) => {
    // Get text before cursor
    const textBeforeCursor = text.substring(0, position);
    const textAfterCursor = text.substring(position);

    // Check if we're typing a LoRA name (but not in the weight part)
    // Match: <lora:name| but NOT <lora:name:weight|
    const loraMatch = textBeforeCursor.match(/<lora:([^:>]*)$/);
    if (loraMatch) {
      const searchTerm = loraMatch[1].toLowerCase();
      const filtered = loras
        .filter(lora => {
          const loraBase = lora.replace(/\.(safetensors|ckpt|pt)$/i, '');
          return loraBase.toLowerCase().includes(searchTerm);
        })
        .slice(0, 10)
        .map(lora => ({
          text: lora,
          type: 'lora' as const,
          displayText: lora.replace(/\.(safetensors|ckpt|pt)$/i, ''),
        }));

      if (filtered.length > 0) {
        setSuggestions(filtered);
        setShowSuggestions(true);
        setSelectedIndex(0);
        return;
      }
    }

    // Check if we're typing an embedding (word boundary)
    const words = textBeforeCursor.split(/\s+/);
    const currentWord = words[words.length - 1];

    // Only show embedding suggestions if we've typed at least 2 characters
    // and we're not inside a lora tag
    if (currentWord.length >= 2 && !textBeforeCursor.match(/<lora:[^>]*$/)) {
      const searchTerm = currentWord.toLowerCase();
      const filtered = embeddings
        .filter(emb => {
          const embBase = emb.replace(/\.(safetensors|pt)$/i, '');
          return embBase.toLowerCase().startsWith(searchTerm);
        })
        .slice(0, 10)
        .map(emb => ({
          text: emb,
          type: 'embedding' as const,
          displayText: emb.replace(/\.(safetensors|pt)$/i, ''),
        }));

      if (filtered.length > 0) {
        setSuggestions(filtered);
        setShowSuggestions(true);
        setSelectedIndex(0);
        return;
      }
    }

    setShowSuggestions(false);
  };

  const insertSuggestion = (suggestion: Suggestion) => {
    if (!textareaRef.current) return;

    const position = textareaRef.current.selectionStart;
    const textBefore = value.substring(0, position);
    let textAfter = value.substring(position);

    let newText = '';
    let newPosition = position;

    if (suggestion.type === 'lora') {
      // Find the <lora: part
      const loraMatch = textBefore.match(/<lora:([^:>]*)$/);
      if (loraMatch) {
        const beforeLora = textBefore.substring(0, textBefore.length - loraMatch[0].length);

        // Check if we're editing an existing tag
        // Remove everything until the closing > (rest of name, weight, closing bracket)
        const afterLoraMatch = textAfter.match(/^[^<>]*>/);
        if (afterLoraMatch) {
          // Remove the old tag remainder
          textAfter = textAfter.substring(afterLoraMatch[0].length);
        }

        const loraText = `<lora:${suggestion.displayText}:0.8>`;
        newText = beforeLora + loraText + textAfter;
        newPosition = beforeLora.length + loraText.length;
      }
    } else {
      // Embedding - replace current word
      const words = textBefore.split(/\s+/);
      const currentWord = words[words.length - 1];
      const beforeWord = textBefore.substring(0, textBefore.length - currentWord.length);
      newText = beforeWord + suggestion.displayText + textAfter;
      newPosition = beforeWord.length + suggestion.displayText.length;
    }

    // Mark that we just inserted a suggestion to prevent retriggering
    justInsertedRef.current = true;

    onChange(newText);
    setShowSuggestions(false);

    // Restore cursor position
    setTimeout(() => {
      if (textareaRef.current) {
        textareaRef.current.selectionStart = newPosition;
        textareaRef.current.selectionEnd = newPosition;
        textareaRef.current.focus();
      }
    }, 0);
  };

  const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
    if (!showSuggestions) return;

    if (e.key === 'ArrowDown') {
      e.preventDefault();
      setSelectedIndex(prev => Math.min(prev + 1, suggestions.length - 1));
    } else if (e.key === 'ArrowUp') {
      e.preventDefault();
      setSelectedIndex(prev => Math.max(prev - 1, 0));
    } else if (e.key === 'Enter' || e.key === 'Tab') {
      if (suggestions.length > 0) {
        e.preventDefault();
        insertSuggestion(suggestions[selectedIndex]);
      }
    } else if (e.key === 'Escape') {
      setShowSuggestions(false);
    }
  };

  const highlightSyntax = (text: string) => {
    if (!text) {
      setHighlighted([]);
      return;
    }

    const parts: React.ReactNode[] = [];
    let lastIndex = 0;

    const loraNames = new Set(
      loras.map(name => name.replace(/\.(safetensors|ckpt|pt)$/i, ''))
    );
    const loraFullNames = new Set(loras);
    const embeddingNames = new Set(
      embeddings.map(name => name.replace(/\.(safetensors|pt)$/i, ''))
    );

    const loraRegex = /<lora:([^:>]+):([^>]+)>/g;
    let match;
    const matches: Array<{ start: number; end: number; type: 'lora' | 'embedding'; text: string; valid: boolean }> = [];

    while ((match = loraRegex.exec(text)) !== null) {
      const loraName = match[1];
      const isValid = loraNames.has(loraName) || loraFullNames.has(loraName);
      matches.push({
        start: match.index,
        end: match.index + match[0].length,
        type: 'lora',
        text: match[0],
        valid: isValid,
      });
    }

    embeddings.forEach(embedding => {
      const embeddingBase = embedding.replace(/\.(safetensors|pt)$/i, '');
      const embeddingRegex = new RegExp(`\\b${escapeRegex(embeddingBase)}\\b`, 'g');
      while ((match = embeddingRegex.exec(text)) !== null) {
        const isInsideLora = matches.some(
          m => m.type === 'lora' && match!.index >= m.start && match!.index < m.end
        );
        if (!isInsideLora) {
          matches.push({
            start: match.index,
            end: match.index + match[0].length,
            type: 'embedding',
            text: match[0],
            valid: true,
          });
        }
      }
    });

    matches.sort((a, b) => a.start - b.start);

    matches.forEach((match, index) => {
      if (match.start > lastIndex) {
        parts.push(
          <span key={`text-${lastIndex}`}>
            {text.substring(lastIndex, match.start)}
          </span>
        );
      }

      const highlightClass = match.type === 'lora'
        ? match.valid
          ? 'bg-purple-500/20 text-purple-700 dark:text-purple-300 font-medium rounded px-0.5'
          : 'bg-red-500/20 text-red-700 dark:text-red-300 font-medium rounded px-0.5'
        : 'bg-blue-500/20 text-blue-700 dark:text-blue-300 font-medium rounded px-0.5';

      parts.push(
        <span key={`highlight-${match.start}`} className={highlightClass} title={match.type === 'lora' ? (match.valid ? 'LoRA' : 'LoRA not found') : 'Embedding'}>
          {match.text}
        </span>
      );

      lastIndex = match.end;
    });

    if (lastIndex < text.length) {
      parts.push(
        <span key={`text-${lastIndex}`}>
          {text.substring(lastIndex)}
        </span>
      );
    }

    setHighlighted(parts);
  };

  const escapeRegex = (str: string) => {
    return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
  };

  return (
    <div className="relative w-full">
      <div
        className="absolute inset-0 pointer-events-none px-3 py-2 text-sm font-mono whitespace-pre-wrap break-words overflow-hidden rounded-md text-foreground"
        style={{
          zIndex: 1,
          WebkitFontSmoothing: 'antialiased',
          MozOsxFontSmoothing: 'grayscale'
        }}
      >
        {highlighted.length > 0 ? highlighted : value}
      </div>
      <textarea
        ref={textareaRef}
        value={value}
        onChange={(e) => {
          onChange(e.target.value);
          setCursorPosition(e.target.selectionStart);
        }}
        onKeyDown={handleKeyDown}
        onClick={(e) => setCursorPosition(e.currentTarget.selectionStart)}
        placeholder={placeholder}
        rows={rows}
        className={cn(
          'prompt-textarea-input flex min-h-[80px] w-full rounded-md border border-input bg-background px-3 py-2 text-sm ring-offset-background placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50 font-mono resize-none relative',
          className
        )}
        style={{
          background: 'transparent',
          color: 'transparent',
          WebkitTextFillColor: 'transparent',
          caretColor: 'hsl(var(--foreground))',
          zIndex: 2,
          textShadow: 'none'
        } as React.CSSProperties & {
          caretColor: string;
        }}
      />

      {/* Autocomplete Suggestions */}
      {showSuggestions && suggestions.length > 0 && (
        <div
          ref={suggestionsRef}
          className="absolute mt-1 w-full max-h-60 overflow-auto rounded-md border border-border bg-popover shadow-lg z-30 pointer-events-auto"
        >
          {suggestions.map((suggestion, index) => (
            <div
              key={`${suggestion.type}-${suggestion.text}`}
              className={cn(
                'px-3 py-2 cursor-pointer text-sm flex items-center justify-between',
                index === selectedIndex
                  ? 'bg-accent text-accent-foreground'
                  : 'hover:bg-accent/50'
              )}
              onClick={() => insertSuggestion(suggestion)}
              onMouseEnter={() => setSelectedIndex(index)}
            >
              <span className="font-mono">{suggestion.displayText}</span>
              <span
                className={cn(
                  'text-xs px-2 py-0.5 rounded',
                  suggestion.type === 'lora'
                    ? 'bg-purple-500/20 text-purple-700 dark:text-purple-300'
                    : 'bg-blue-500/20 text-blue-700 dark:text-blue-300'
                )}
              >
                {suggestion.type}
              </span>
            </div>
          ))}
        </div>
      )}
    </div>
  );
}