searchBar.tsx

'use client';

import { useClickListener } from "@/hooks/useClickListener";
import { SearchQueryParams } from "@/lib/types";
import { cn, createPathWithQueryParams } from "@/lib/utils";
import {
    cursorCharLeft,
    cursorCharRight,
    cursorDocEnd,
    cursorDocStart,
    cursorLineBoundaryBackward,
    cursorLineBoundaryForward,
    deleteCharBackward,
    deleteCharForward,
    deleteGroupBackward,
    deleteGroupForward,
    deleteLineBoundaryBackward,
    deleteLineBoundaryForward,
    history,
    historyKeymap,
    selectAll,
    selectCharLeft,
    selectCharRight,
    selectDocEnd,
    selectDocStart,
    selectLineBoundaryBackward,
    selectLineBoundaryForward
} from "@codemirror/commands";
import CodeMirror, { Annotation, EditorView, KeyBinding, keymap, ReactCodeMirrorRef } from "@uiw/react-codemirror";
import { cva } from "class-variance-authority";
import { useRouter } from "next/navigation";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { useHotkeys } from 'react-hotkeys-hook';
import { SearchSuggestionsBox } from "./searchSuggestionsBox";
import { useSuggestionsData } from "./useSuggestionsData";
import { zoekt } from "./zoektLanguageExtension";
import { CounterClockwiseClockIcon } from "@radix-ui/react-icons";
import { useSuggestionModeAndQuery } from "./useSuggestionModeAndQuery";
import { Separator } from "@/components/ui/separator";
import { Tooltip, TooltipTrigger, TooltipContent } from "@/components/ui/tooltip";
import { Toggle } from "@/components/ui/toggle";
import { useDomain } from "@/hooks/useDomain";
import React from "react";
import Link from "next/link";
import { CaseSensitiveIcon, RegexIcon, Wand2Icon } from "lucide-react";
import { SearchAssistBox } from "./searchAssistBox";
import useCaptureEvent from "@/hooks/useCaptureEvent";
import { useCodeMirrorTheme } from "@/hooks/useCodeMirrorTheme";

const LANGUAGE_MODEL_DOCS_URL = "https://docs.sourcebot.dev/docs/configuration/language-model-providers"; 

interface SearchBarProps {
    className?: string;
    size?: "default" | "sm";
    defaults?: {
        isRegexEnabled?: boolean;
        isCaseSensitivityEnabled?: boolean;
        query?: string;
    }
    autoFocus?: boolean;
    isSearchAssistSupported: boolean;
}

const searchBarKeymap: readonly KeyBinding[] = ([
    { key: "ArrowLeft", run: cursorCharLeft, shift: selectCharLeft, preventDefault: true },
    { key: "ArrowRight", run: cursorCharRight, shift: selectCharRight, preventDefault: true },

    { key: "Home", run: cursorLineBoundaryBackward, shift: selectLineBoundaryBackward, preventDefault: true },
    { key: "Mod-Home", run: cursorDocStart, shift: selectDocStart },

    { key: "End", run: cursorLineBoundaryForward, shift: selectLineBoundaryForward, preventDefault: true },
    { key: "Mod-End", run: cursorDocEnd, shift: selectDocEnd },

    { key: "Mod-a", run: selectAll },

    { key: "Backspace", run: deleteCharBackward, shift: deleteCharBackward },
    { key: "Delete", run: deleteCharForward },
    { key: "Mod-Backspace", mac: "Alt-Backspace", run: deleteGroupBackward },
    { key: "Mod-Delete", mac: "Alt-Delete", run: deleteGroupForward },
    { mac: "Mod-Backspace", run: deleteLineBoundaryBackward },
    { mac: "Mod-Delete", run: deleteLineBoundaryForward }
] as KeyBinding[]).concat(historyKeymap);

const searchBarContainerVariants = cva(
    "search-bar-container flex items-center justify-center py-0.5 px-2 border rounded-md relative",
    {
        variants: {
            size: {
                default: "min-h-10",
                sm: "min-h-8"
            }
        },
        defaultVariants: {
            size: "default",
        }
    }
);

export const SearchBar = ({
    className,
    size,
    autoFocus,
    defaults: {
        isRegexEnabled: defaultIsRegexEnabled = false,
        isCaseSensitivityEnabled: defaultIsCaseSensitivityEnabled = false,
        query: defaultQuery = "",
    } = {},
    isSearchAssistSupported,
}: SearchBarProps) => {
    const router = useRouter();
    const domain = useDomain();
    const captureEvent = useCaptureEvent();
    const suggestionBoxRef = useRef<HTMLDivElement>(null);
    const editorRef = useRef<ReactCodeMirrorRef>(null);
    const [cursorPosition, setCursorPosition] = useState(0);
    const [activePanel, setActivePanel] = useState<'suggestions' | 'searchAssist'>();
    const isSuggestionsEnabled = activePanel === 'suggestions';
    const isSearchAssistEnabled = activePanel === 'searchAssist';
    const [isSuggestionsBoxFocused, setIsSuggestionsBoxFocused] = useState(false);
    const [isHistorySearchEnabled, setIsHistorySearchEnabled] = useState(false);
    const [isRegexEnabled, setIsRegexEnabled] = useState(defaultIsRegexEnabled);
    const [isCaseSensitivityEnabled, setIsCaseSensitivityEnabled] = useState(defaultIsCaseSensitivityEnabled);

    const focusEditor = useCallback(() => editorRef.current?.view?.focus(), []);
    const focusSuggestionsBox = useCallback(() => suggestionBoxRef.current?.focus(), []);

    const [_query, setQuery] = useState(defaultQuery);
    const query = useMemo(() => {
        // Replace any newlines with spaces to handle
        // copy & pasting text with newlines.
        return _query.replaceAll(/\n/g, " ");
    }, [_query]);

    // When the user navigates backwards/forwards while on the
    // search page (causing the `query` search param to change),
    // we want to update what query is displayed in the search bar.
    useEffect(() => {
        if (defaultQuery) {
            setQuery(defaultQuery);
        }
    }, [defaultQuery])


    const { suggestionMode, suggestionQuery } = useSuggestionModeAndQuery({
        isSuggestionsEnabled,
        isHistorySearchEnabled,
        cursorPosition,
        query,
    });

    const suggestionData = useSuggestionsData({
        suggestionMode,
        suggestionQuery,
    });

    const theme = useCodeMirrorTheme();

    const extensions = useMemo(() => {
        return [
            zoekt(),
            keymap.of(searchBarKeymap),
            history(),
            EditorView.lineWrapping,
            EditorView.updateListener.of(update => {
                if (update.selectionSet) {
                    const selection = update.state.selection.main;
                    if (selection.empty) {
                        setCursorPosition(selection.anchor);
                    }
                }
            })
        ];
    }, []);

    // Hotkey to focus the search bar.
    useHotkeys('/', (event) => {
        event.preventDefault();
        focusEditor();
        setActivePanel('suggestions');
        if (editorRef.current?.view) {
            cursorDocEnd({
                state: editorRef.current.view.state,
                dispatch: editorRef.current.view.dispatch,
            });
        }
    });

    // Collapse the suggestions box if the user clicks outside of the search bar container.
    useClickListener('.search-bar-container', (isElementClicked) => {
        if (!isElementClicked) {
            setActivePanel(undefined);
        } else {
            setActivePanel(prev => prev ?? 'suggestions');
        }
    });

    const onSubmit = useCallback((query: string) => {
        setActivePanel(undefined);
        setIsHistorySearchEnabled(false);

        const url = createPathWithQueryParams(`/${domain}/search`,
            [SearchQueryParams.query, query],
            [SearchQueryParams.isRegexEnabled, isRegexEnabled ? "true" : null],
            [SearchQueryParams.isCaseSensitivityEnabled, isCaseSensitivityEnabled ? "true" : null],
        );
        router.push(url);
    }, [domain, router, isRegexEnabled, isCaseSensitivityEnabled]);

    return (
        <div
            className={cn(searchBarContainerVariants({ size, className }))}
            onKeyDown={(e) => {
                if (e.key === 'Enter') {
                    e.preventDefault();
                    if (activePanel !== 'searchAssist') {
                        setActivePanel(undefined);
                        onSubmit(query);
                    }
                }

                if (e.key === 'Escape') {
                    e.preventDefault();
                    setActivePanel(undefined);
                }

                if (e.key === 'ArrowDown' && !isSearchAssistEnabled) {
                    e.preventDefault();
                    setActivePanel('suggestions');
                    focusSuggestionsBox();
                }

                if (e.key === 'ArrowUp') {
                    e.preventDefault();
                }
            }}
        >
            <div className="flex flex-row items-center gap-1">
                <SearchBarButton
                    isToggled={isHistorySearchEnabled}
                    onClick={() => {
                        setQuery("");
                        setIsHistorySearchEnabled(!isHistorySearchEnabled);
                        setActivePanel('suggestions');
                        focusEditor();
                    }}
                    tooltip="Search history"
                    icon={CounterClockwiseClockIcon}
                />
                <SearchBarButton
                    isToggled={isSearchAssistEnabled}
                    onClick={() => {
                        setQuery("");
                        setIsHistorySearchEnabled(false);
                        setActivePanel(prev => {
                            const next = prev === 'searchAssist' ? undefined : 'searchAssist';
                            if (next === 'searchAssist') {
                                captureEvent('wa_search_assist_opened', {});
                            }
                            return next;
                        });
                        focusEditor();
                    }}
                    tooltip="AI search assist"
                    icon={Wand2Icon}
                    preventBlurOnClick
                    disabled={!isSearchAssistSupported}
                    disabledTooltip={
                        <span>
                            AI search assist requires a language model to be configured.{" "}
                            <Link href={LANGUAGE_MODEL_DOCS_URL} target="_blank" className="underline">
                                Learn more
                            </Link>.
                        </span>
                    }
                />
            </div>
            <Separator
                className="mx-1 h-6"
                orientation="vertical"
            />
            <CodeMirror
                ref={editorRef}
                className="w-full"
                placeholder={isHistorySearchEnabled ? "Filter history..." : "Search (/) through repos..."}
                value={query}
                onChange={(value) => {
                    setQuery(value);
                    // Whenever the user types, we want to re-enable
                    // the suggestions box.
                    setActivePanel('suggestions');
                }}
                theme={theme}
                basicSetup={false}
                extensions={extensions}
                indentWithTab={false}
                autoFocus={autoFocus ?? false}
            />
            <div className="flex flex-row items-center gap-1 ml-1">
                <SearchBarButton
                    isToggled={isCaseSensitivityEnabled}
                    onClick={() => setIsCaseSensitivityEnabled(!isCaseSensitivityEnabled)}
                    tooltip={`${isCaseSensitivityEnabled ? "Disable" : "Enable"} case sensitivity`}
                    icon={CaseSensitiveIcon}

                />
                <SearchBarButton
                    isToggled={isRegexEnabled}
                    onClick={() => setIsRegexEnabled(!isRegexEnabled)}
                    tooltip={`${isRegexEnabled ? "Disable" : "Enable"} regular expressions`}
                    icon={RegexIcon}

                />
            </div>
            <SearchAssistBox
                className={size === "sm" ? "top-7" : "top-9"}
                isEnabled={isSearchAssistEnabled}
                onBlur={() => {
                    setActivePanel(undefined);
                }}
                onQueryGenerated={(translatedQuery: string) => {
                    setQuery(translatedQuery);
                    editorRef.current?.view?.dispatch({
                        changes: { from: 0, to: editorRef.current.view.state.doc.length, insert: translatedQuery },
                        selection: { anchor: translatedQuery.length },
                    });
                    setActivePanel(undefined);
                    focusEditor();
                    // Always enable regex and case sensitivity when using search assist.
                    setIsRegexEnabled(true);
                    setIsCaseSensitivityEnabled(true);
                }}
            />
            <SearchSuggestionsBox
                ref={suggestionBoxRef}
                className={size === "sm" ? "top-9" : "top-12"}
                query={query}
                suggestionQuery={suggestionQuery}
                suggestionMode={suggestionMode}
                onCompletion={(newQuery: string, newCursorPosition: number, autoSubmit = false) => {
                    setQuery(newQuery);

                    // Move the cursor to it's new position.
                    // @note : normally, react-codemirror handles syncing `query`
                    // and the document state, but this happens on re-render. Since
                    // we want to move the cursor before the component re-renders,
                    // we manually update the document state inline.
                    editorRef.current?.view?.dispatch({
                        changes: { from: 0, to: query.length, insert: newQuery },
                        annotations: [Annotation.define<boolean>().of(true)],
                    });

                    editorRef.current?.view?.dispatch({
                        selection: { anchor: newCursorPosition, head: newCursorPosition },
                    });

                    // Re-focus the editor since suggestions cause focus to be lost (both click & keyboard)
                    editorRef.current?.view?.focus();

                    if (autoSubmit) {
                        onSubmit(newQuery);
                    }
                }}
                isEnabled={isSuggestionsEnabled}
                onReturnFocus={() => {
                    focusEditor();
                }}
                isFocused={isSuggestionsBoxFocused}
                onFocus={() => {
                    setIsSuggestionsBoxFocused(document.activeElement === suggestionBoxRef.current);
                }}
                onBlur={() => {
                    setIsSuggestionsBoxFocused(document.activeElement === suggestionBoxRef.current);
                }}
                cursorPosition={cursorPosition}
                {...suggestionData}
            />
        </div>
    )
}

const SearchBarButton = ({
    isToggled,
    onClick,
    tooltip,
    icon: Icon,
    preventBlurOnClick = false,
    disabled = false,
    disabledTooltip,
}: {
    isToggled: boolean,
    onClick: () => void,
    tooltip: React.ReactNode,
    icon: React.ElementType,
    preventBlurOnClick?: boolean,
    disabled?: boolean,
    disabledTooltip?: React.ReactNode,
}) => {
    return (
        <Tooltip>
            <TooltipTrigger asChild={true}>
                {/* @see : https://github.com/shadcn-ui/ui/issues/1988#issuecomment-1980597269 */}
                <div>
                    <Toggle
                        pressed={isToggled}
                        className="h-6 w-6 min-w-6 px-0 p-1 cursor-pointer"
                        onClick={onClick}
                        onMouseDown={preventBlurOnClick ? (e) => e.preventDefault() : undefined}
                        disabled={disabled}
                    >
                        <Icon className="w-4 h-4" />
                    </Toggle>
                </div>
            </TooltipTrigger>
            <TooltipContent side="bottom">
                {disabled && disabledTooltip ? disabledTooltip : tooltip}
            </TooltipContent>
        </Tooltip>
    )
}