"use client";

import { ChatInput } from "@llamaindex/chat-ui";
import { FolderOpen, Check } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
    Dialog,
    DialogContent,
    DialogDescription,
    DialogTitle,
    DialogTrigger,
    DialogFooter,
} from "@/components/ui/dialog";
import {
    Select,
    SelectContent,
    SelectItem,
    SelectTrigger,
    SelectValue,
} from "@/components/ui/select";
import { Suspense, useState, useCallback } from "react";
import { useParams } from "next/navigation";
import { useDocuments, Document } from "@/hooks/use-documents";
import { DocumentsDataTable } from "@/components/chat_v2/DocumentsDataTable";
import { useSearchSourceConnectors } from "@/hooks/useSearchSourceConnectors";
import {
    getConnectorIcon,
    ConnectorButton as ConnectorButtonComponent,
} from "@/components/chat/ConnectorComponents";
import { ResearchMode } from "@/components/chat";
import React from "react";

const DocumentSelector = React.memo(
    ({
        onSelectionChange,
        selectedDocuments = [],
    }: {
        onSelectionChange?: (documents: Document[]) => void;
        selectedDocuments?: Document[];
    }) => {
        const { search_space_id } = useParams();
        const [isOpen, setIsOpen] = useState(false);

        const { documents, loading, isLoaded, fetchDocuments } = useDocuments(
            Number(search_space_id),
            true
        );

        const handleOpenChange = useCallback(
            (open: boolean) => {
                setIsOpen(open);
                if (open && !isLoaded) {
                    fetchDocuments();
                }
            },
            [fetchDocuments, isLoaded]
        );

        const handleSelectionChange = useCallback(
            (documents: Document[]) => {
                onSelectionChange?.(documents);
            },
            [onSelectionChange]
        );

        const handleDone = useCallback(() => {
            setIsOpen(false);
        }, [selectedDocuments]);

        const selectedCount = selectedDocuments.length;

        return (
            <Dialog open={isOpen} onOpenChange={handleOpenChange}>
                <DialogTrigger asChild>
                    <Button variant="outline" className="relative">
                        <FolderOpen className="w-4 h-4" />
                        {selectedCount > 0 && (
                            <span className="absolute -top-2 -right-2 h-5 w-5 rounded-full bg-primary text-primary-foreground text-xs flex items-center justify-center">
                                {selectedCount}
                            </span>
                        )}
                    </Button>
                </DialogTrigger>

                <DialogContent className="max-w-[95vw] md:max-w-5xl h-[90vh] md:h-[85vh] p-0 flex flex-col">
                    <div className="flex flex-col h-full">
                        <div className="px-4 md:px-6 py-4 border-b flex-shrink-0">
                            <DialogTitle className="text-lg md:text-xl">
                                Select Documents
                            </DialogTitle>
                            <DialogDescription className="mt-1 text-sm">
                                Choose documents to include in your research
                                context
                            </DialogDescription>
                        </div>

                        <div className="flex-1 min-h-0 p-4 md:p-6">
                            {loading ? (
                                <div className="flex items-center justify-center h-full">
                                    <div className="text-center space-y-2">
                                        <div className="animate-spin h-8 w-8 border-2 border-primary border-t-transparent rounded-full mx-auto" />
                                        <p className="text-sm text-muted-foreground">
                                            Loading documents...
                                        </p>
                                    </div>
                                </div>
                            ) : isLoaded ? (
                                <DocumentsDataTable
                                    documents={documents}
                                    onSelectionChange={handleSelectionChange}
                                    onDone={handleDone}
                                    initialSelectedDocuments={selectedDocuments}
                                />
                            ) : null}
                        </div>
                    </div>
                </DialogContent>
            </Dialog>
        );
    }
);

const ConnectorSelector = React.memo(
    ({
        onSelectionChange,
        selectedConnectors = [],
    }: {
        onSelectionChange?: (connectorTypes: string[]) => void;
        selectedConnectors?: string[];
    }) => {
        const [isOpen, setIsOpen] = useState(false);

        const { connectorSourceItems, isLoading, isLoaded, fetchConnectors } =
            useSearchSourceConnectors();

        const handleOpenChange = useCallback(
            (open: boolean) => {
                setIsOpen(open);
                if (open && !isLoaded) {
                    fetchConnectors();
                }
            },
            [fetchConnectors, isLoaded]
        );

        const handleConnectorToggle = useCallback(
            (connectorType: string) => {
                const isSelected = selectedConnectors.includes(connectorType);
                const newSelection = isSelected
                    ? selectedConnectors.filter(
                          (type) => type !== connectorType
                      )
                    : [...selectedConnectors, connectorType];
                onSelectionChange?.(newSelection);
            },
            [selectedConnectors, onSelectionChange]
        );

        const handleSelectAll = useCallback(() => {
            onSelectionChange?.(connectorSourceItems.map((c) => c.type));
        }, [connectorSourceItems, onSelectionChange]);

        const handleClearAll = useCallback(() => {
            onSelectionChange?.([]);
        }, [onSelectionChange]);

        return (
            <Dialog open={isOpen} onOpenChange={handleOpenChange}>
                <DialogTrigger asChild>
                    <ConnectorButtonComponent
                        selectedConnectors={selectedConnectors}
                        onClick={() => setIsOpen(true)}
                        connectorSources={connectorSourceItems}
                    />
                </DialogTrigger>

                <DialogContent className="sm:max-w-md">
                    <DialogTitle>Select Connectors</DialogTitle>
                    <DialogDescription>
                        Choose which data sources to include in your research
                    </DialogDescription>

                    {/* Connector selection grid */}
                    <div className="grid grid-cols-2 gap-4 py-4">
                        {isLoading ? (
                            <div className="col-span-2 flex justify-center py-4">
                                <div className="animate-spin h-6 w-6 border-2 border-primary border-t-transparent rounded-full" />
                            </div>
                        ) : (
                            connectorSourceItems.map((connector) => {
                                const isSelected = selectedConnectors.includes(
                                    connector.type
                                );

                                return (
                                    <div
                                        key={connector.id}
                                        className={`flex items-center gap-2 p-2 rounded-md border cursor-pointer transition-colors ${
                                            isSelected
                                                ? "border-primary bg-primary/10"
                                                : "border-border hover:border-primary/50 hover:bg-muted"
                                        }`}
                                        onClick={() =>
                                            handleConnectorToggle(
                                                connector.type
                                            )
                                        }
                                        role="checkbox"
                                        aria-checked={isSelected}
                                        tabIndex={0}
                                    >
                                        <div className="flex-shrink-0 w-6 h-6 flex items-center justify-center rounded-full bg-muted">
                                            {getConnectorIcon(connector.type)}
                                        </div>
                                        <span className="flex-1 text-sm font-medium">
                                            {connector.name}
                                        </span>
                                        {isSelected && (
                                            <Check className="h-4 w-4 text-primary" />
                                        )}
                                    </div>
                                );
                            })
                        )}
                    </div>

                    <DialogFooter className="flex justify-between items-center">
                        <div className="flex gap-2">
                            <Button variant="outline" onClick={handleClearAll}>
                                Clear All
                            </Button>
                            <Button onClick={handleSelectAll}>
                                Select All
                            </Button>
                        </div>
                    </DialogFooter>
                </DialogContent>
            </Dialog>
        );
    }
);

const SearchModeSelector = ({
    searchMode,
    onSearchModeChange,
}: {
    searchMode?: "DOCUMENTS" | "CHUNKS";
    onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
}) => {
    return (
        <div className="flex items-center gap-1 sm:gap-2">
            <span className="text-xs text-muted-foreground hidden sm:block">
                Scope:
            </span>
            <div className="flex rounded-md border">
                <Button
                    variant={searchMode === "DOCUMENTS" ? "default" : "ghost"}
                    size="sm"
                    className="rounded-r-none border-r h-8 px-2 sm:px-3 text-xs"
                    onClick={() => onSearchModeChange?.("DOCUMENTS")}
                >
                    <span className="hidden sm:inline">Documents</span>
                    <span className="sm:hidden">Docs</span>
                </Button>
                <Button
                    variant={searchMode === "CHUNKS" ? "default" : "ghost"}
                    size="sm"
                    className="rounded-l-none h-8 px-2 sm:px-3 text-xs"
                    onClick={() => onSearchModeChange?.("CHUNKS")}
                >
                    Chunks
                </Button>
            </div>
        </div>
    );
};

const ResearchModeSelector = ({
    researchMode,
    onResearchModeChange,
}: {
    researchMode?: ResearchMode;
    onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
    return (
        <div className="flex items-center gap-1 sm:gap-2">
            <span className="text-xs text-muted-foreground hidden sm:block">
                Mode:
            </span>
            <Select
                value={researchMode}
                onValueChange={(value) =>
                    onResearchModeChange?.(value as ResearchMode)
                }
            >
                <SelectTrigger className="w-auto min-w-[80px] sm:min-w-[120px] h-8 text-xs">
                    <SelectValue placeholder="Mode" />
                </SelectTrigger>
                <SelectContent>
                    <SelectItem value="QNA">Q&A</SelectItem>
                    <SelectItem value="REPORT_GENERAL">
                        <span className="hidden sm:inline">General Report</span>
                        <span className="sm:hidden">General</span>
                    </SelectItem>
                    <SelectItem value="REPORT_DEEP">
                        <span className="hidden sm:inline">Deep Report</span>
                        <span className="sm:hidden">Deep</span>
                    </SelectItem>
                    <SelectItem value="REPORT_DEEPER">
                        <span className="hidden sm:inline">Deeper Report</span>
                        <span className="sm:hidden">Deeper</span>
                    </SelectItem>
                </SelectContent>
            </Select>
        </div>
    );
};

const CustomChatInputOptions = ({
    onDocumentSelectionChange,
    selectedDocuments,
    onConnectorSelectionChange,
    selectedConnectors,
    searchMode,
    onSearchModeChange,
    researchMode,
    onResearchModeChange,
}: {
    onDocumentSelectionChange?: (documents: Document[]) => void;
    selectedDocuments?: Document[];
    onConnectorSelectionChange?: (connectorTypes: string[]) => void;
    selectedConnectors?: string[];
    searchMode?: "DOCUMENTS" | "CHUNKS";
    onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
    researchMode?: ResearchMode;
    onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
    return (
        <div className="flex flex-wrap gap-2 sm:gap-3 items-center justify-start">
            <Suspense fallback={<div>Loading...</div>}>
                <DocumentSelector
                    onSelectionChange={onDocumentSelectionChange}
                    selectedDocuments={selectedDocuments}
                />
            </Suspense>
            <Suspense fallback={<div>Loading...</div>}>
                <ConnectorSelector
                    onSelectionChange={onConnectorSelectionChange}
                    selectedConnectors={selectedConnectors}
                />
            </Suspense>
            <SearchModeSelector
                searchMode={searchMode}
                onSearchModeChange={onSearchModeChange}
            />
            <ResearchModeSelector
                researchMode={researchMode}
                onResearchModeChange={onResearchModeChange}
            />
        </div>
    );
};

export const CustomChatInput = ({
    onDocumentSelectionChange,
    selectedDocuments,
    onConnectorSelectionChange,
    selectedConnectors,
    searchMode,
    onSearchModeChange,
    researchMode,
    onResearchModeChange,
}: {
    onDocumentSelectionChange?: (documents: Document[]) => void;
    selectedDocuments?: Document[];
    onConnectorSelectionChange?: (connectorTypes: string[]) => void;
    selectedConnectors?: string[];
    searchMode?: "DOCUMENTS" | "CHUNKS";
    onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
    researchMode?: ResearchMode;
    onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
    return (
        <ChatInput>
            <ChatInput.Form className="flex gap-2">
                <ChatInput.Field className="flex-1" />
                <ChatInput.Submit />
            </ChatInput.Form>
            <CustomChatInputOptions
                onDocumentSelectionChange={onDocumentSelectionChange}
                selectedDocuments={selectedDocuments}
                onConnectorSelectionChange={onConnectorSelectionChange}
                selectedConnectors={selectedConnectors}
                searchMode={searchMode}
                onSearchModeChange={onSearchModeChange}
                researchMode={researchMode}
                onResearchModeChange={onResearchModeChange}
            />
        </ChatInput>
    );
};