/* Copyright 2026 Marimo. All rights reserved. */

import { CommandList } from "cmdk";
import {
  CopyIcon,
  DownloadIcon,
  FolderIcon,
  HardDriveIcon,
  HelpCircleIcon,
  LoaderCircle,
  PlusIcon,
  ViewIcon,
  XIcon,
} from "lucide-react";
import React, { useCallback, useState } from "react";
import { useLocale } from "react-aria";
import { EngineVariable } from "@/components/databases/engine-variable";
import { useAddCodeToNewCell } from "@/components/editor/cell/useAddCell";
import { PanelEmptyState } from "@/components/editor/chrome/panels/empty-state";
import { AddConnectionDialog } from "@/components/editor/connections/add-connection-dialog";
import {
  FILE_ICON_COLOR,
  renderFileIcon,
} from "@/components/editor/file-tree/file-icons";
import {
  MENU_ITEM_ICON_CLASS,
  MoreActionsButton,
  RefreshIconButton,
  TreeChevron,
} from "@/components/editor/file-tree/tree-actions";
import { Command, CommandInput, CommandItem } from "@/components/ui/command";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuSeparator,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Tooltip } from "@/components/ui/tooltip";
import { toast } from "@/components/ui/use-toast";
import { DownloadStorage } from "@/core/storage/request-registry";
import {
  useStorage,
  useStorageActions,
  useStorageEntries,
} from "@/core/storage/state";
import type {
  StorageEntry,
  StorageNamespace,
  StoragePathKey,
} from "@/core/storage/types";
import { storagePathKey } from "@/core/storage/types";
import { cn } from "@/utils/cn";
import { copyToClipboard } from "@/utils/copy";
import { downloadByURL } from "@/utils/download";
import { formatBytes } from "@/utils/formatting";
import { Logger } from "@/utils/Logger";
import { ErrorState } from "../datasources/components";
import { Button } from "../ui/button";
import { ProtocolIcon } from "./components";
import { StorageFileViewer } from "./storage-file-viewer";
import { STORAGE_SNIPPETS } from "./storage-snippets";

interface OpenFileInfo {
  entry: StorageEntry;
  namespace: string;
  protocol: string;
  backendType: StorageNamespace["backendType"];
}

// Pixels per depth level. Applied as paddingLeft on each full-width item
// so the selection highlight still spans the entire panel.
const INDENT_PX = 16;

function indentStyle(depth: number): React.CSSProperties {
  return { paddingLeft: depth * INDENT_PX };
}

function formatDate(timestamp: number, locale: string): string {
  const date = new Date(timestamp * 1000);
  return date.toLocaleDateString(locale, {
    month: "short",
    day: "numeric",
    hour: "2-digit",
    minute: "2-digit",
  });
}

/** Extract display name from a full path (e.g., "folder/subfolder/" -> "subfolder") */
function displayName(path: string): string {
  // Remove trailing slash
  const trimmed = path.endsWith("/") ? path.slice(0, -1) : path;
  const parts = trimmed.split("/");
  return parts[parts.length - 1] || trimmed;
}

/**
 * Recursively check whether an entry (or any of its loaded descendants)
 * matches the search query.
 */
function entryMatchesSearch(
  entry: StorageEntry,
  namespace: string,
  searchValue: string,
  entriesByPath: ReadonlyMap<StoragePathKey, StorageEntry[]>,
): boolean {
  const query = searchValue.toLowerCase();

  if (displayName(entry.path).toLowerCase().includes(query)) {
    return true;
  }

  // For directories, check loaded children recursively
  if (entry.kind === "directory") {
    const children = entriesByPath.get(storagePathKey(namespace, entry.path));
    if (children) {
      return children.some((child) =>
        entryMatchesSearch(child, namespace, searchValue, entriesByPath),
      );
    }
  }

  return false;
}

/**
 * Filter entries to those matching the search (or having loaded descendants
 * that match). Returns all entries when there is no active search.
 */
function filterEntries(
  entries: StorageEntry[],
  namespace: string,
  searchValue: string,
  entriesByPath: ReadonlyMap<StoragePathKey, StorageEntry[]>,
): StorageEntry[] {
  if (!searchValue.trim()) {
    return entries;
  }
  return entries.filter((entry) =>
    entryMatchesSearch(entry, namespace, searchValue, entriesByPath),
  );
}

/**
 * Lazily loaded children of a directory entry.
 * Caches fetched entries in the Jotai store so re-expanding doesn't re-fetch.
 */
const StorageEntryChildren: React.FC<{
  namespace: string;
  protocol: string;
  rootPath: string;
  backendType: StorageNamespace["backendType"];
  prefix: string;
  depth: number;
  locale: string;
  searchValue: string;
  onOpenFile: (info: OpenFileInfo) => void;
}> = ({
  namespace,
  protocol,
  rootPath,
  backendType,
  prefix,
  depth,
  locale,
  searchValue,
  onOpenFile,
}) => {
  const { entriesByPath } = useStorage();
  const {
    entries: children,
    isPending,
    error,
  } = useStorageEntries(namespace, prefix);

  if (isPending) {
    return (
      <div
        className="flex items-center gap-1.5 py-1 text-xs text-muted-foreground"
        style={indentStyle(depth)}
      >
        <LoaderCircle className="h-3 w-3 animate-spin" />
        Loading...
      </div>
    );
  }

  if (error) {
    return (
      <div className="py-1 text-xs text-destructive" style={indentStyle(depth)}>
        Failed to load: {error.message}
      </div>
    );
  }

  if (children.length === 0) {
    return (
      <div
        className="py-1 text-xs text-muted-foreground italic"
        style={indentStyle(depth)}
      >
        Empty
      </div>
    );
  }

  const filtered = filterEntries(
    children,
    namespace,
    searchValue,
    entriesByPath,
  );

  return (
    <>
      {filtered.map((child) => (
        <StorageEntryRow
          key={child.path}
          entry={child}
          namespace={namespace}
          protocol={protocol}
          rootPath={rootPath}
          backendType={backendType}
          depth={depth}
          locale={locale}
          searchValue={searchValue}
          onOpenFile={onOpenFile}
        />
      ))}
    </>
  );
};

const StorageEntryRow: React.FC<{
  entry: StorageEntry;
  namespace: string;
  protocol: string;
  rootPath: string;
  backendType: StorageNamespace["backendType"];
  depth: number;
  locale: string;
  searchValue: string;
  onOpenFile: (info: OpenFileInfo) => void;
}> = ({
  entry,
  namespace,
  protocol,
  rootPath,
  backendType,
  depth,
  locale,
  searchValue,
  onOpenFile,
}) => {
  const [isExpanded, setIsExpanded] = useState(false);
  const { entriesByPath } = useStorage();
  const addCodeToNewCell = useAddCodeToNewCell();
  const isDir = entry.kind === "directory";
  const name = displayName(entry.path);
  const hasSearch = !!searchValue.trim();

  const selfMatches =
    isDir &&
    hasSearch &&
    name.toLowerCase().includes(searchValue.trim().toLowerCase());

  // During a search, auto-expand directories whose loaded descendants match
  const hasMatchingDescendants =
    isDir &&
    hasSearch &&
    !!entriesByPath
      .get(storagePathKey(namespace, entry.path))
      ?.some((child) =>
        entryMatchesSearch(child, namespace, searchValue, entriesByPath),
      );

  // Folder is shown expanded by manual toggle OR by search auto-expand
  const effectiveExpanded = isExpanded || hasMatchingDescendants;

  const handleDownload = useCallback(async () => {
    try {
      const result = await DownloadStorage.request({
        namespace,
        path: entry.path,
      });
      if (result.error) {
        toast({
          title: "Download failed",
          description: result.error,
          variant: "danger",
        });
        return;
      }
      if (result.url) {
        downloadByURL(result.url, result.filename ?? "download");
      }
    } catch (error) {
      Logger.error("Failed to download storage entry", error);
      toast({
        title: "Download failed",
        description: String(error),
        variant: "danger",
      });
    }
  }, [namespace, entry.path]);

  return (
    <>
      <CommandItem
        className={cn(
          "text-xs flex items-center gap-1.5 cursor-pointer rounded-none group h-6.5",
          isDir && "font-medium",
        )}
        style={indentStyle(depth)}
        value={`${namespace}:${entry.path}`}
        onSelect={() => {
          if (isDir) {
            setIsExpanded(!effectiveExpanded);
          } else {
            onOpenFile({ entry, namespace, protocol, backendType });
          }
        }}
      >
        {isDir ? (
          <TreeChevron isExpanded={effectiveExpanded} className="h-3 w-3" />
        ) : (
          <span className="w-3 shrink-0" />
        )}
        {isDir ? (
          <FolderIcon
            className={cn("h-3.5 w-3.5 shrink-0", FILE_ICON_COLOR.directory)}
          />
        ) : (
          renderFileIcon(name)
        )}
        <span className="truncate flex-1 text-left">{name}</span>
        <div className="flex items-center">
          {entry.size > 0 && (
            <span className="text-[10px] text-muted-foreground pr-2 opacity-0 group-hover:opacity-100 transition-opacity tabular-nums">
              {formatBytes(entry.size, locale)}
            </span>
          )}
          {entry.lastModified != null && (
            <Tooltip
              content={`Last modified: ${new Date(entry.lastModified * 1000).toLocaleString()}`}
            >
              <span className="text-[10px] text-muted-foreground pr-1 opacity-0 group-hover:opacity-100 transition-opacity whitespace-nowrap">
                {formatDate(entry.lastModified, locale)}
              </span>
            </Tooltip>
          )}
          <DropdownMenu>
            <DropdownMenuTrigger asChild={true}>
              <MoreActionsButton
                iconClassName="h-3 w-3"
                onClick={(e) => e.stopPropagation()}
              />
            </DropdownMenuTrigger>
            <DropdownMenuContent
              align="end"
              onClick={(e) => e.stopPropagation()}
              onCloseAutoFocus={(e) => e.preventDefault()}
            >
              {!isDir && (
                <DropdownMenuItem
                  onSelect={() =>
                    onOpenFile({ entry, namespace, protocol, backendType })
                  }
                >
                  <ViewIcon className={MENU_ITEM_ICON_CLASS} />
                  View
                </DropdownMenuItem>
              )}
              <DropdownMenuItem
                onSelect={async () => {
                  await copyToClipboard(entry.path);
                  toast({ title: "Copied to clipboard" });
                }}
              >
                <CopyIcon className={MENU_ITEM_ICON_CLASS} />
                Copy path
              </DropdownMenuItem>
              {!isDir && (
                <DropdownMenuItem onSelect={() => handleDownload()}>
                  <DownloadIcon className={MENU_ITEM_ICON_CLASS} />
                  Download
                </DropdownMenuItem>
              )}
              <DropdownMenuSeparator />
              {STORAGE_SNIPPETS.map((snippet) => {
                const code = snippet.getCode({
                  variableName: namespace,
                  protocol,
                  entry,
                  backendType,
                });
                if (code === null) {
                  return null;
                }
                const Icon = snippet.icon;
                return (
                  <DropdownMenuItem
                    key={snippet.id}
                    onSelect={() => addCodeToNewCell(code)}
                  >
                    <Icon className={MENU_ITEM_ICON_CLASS} />
                    {snippet.label}
                  </DropdownMenuItem>
                );
              })}
            </DropdownMenuContent>
          </DropdownMenu>
        </div>
      </CommandItem>
      {isDir && effectiveExpanded && (
        <StorageEntryChildren
          namespace={namespace}
          protocol={protocol}
          rootPath={rootPath}
          backendType={backendType}
          prefix={entry.path}
          depth={depth + 1}
          locale={locale}
          searchValue={selfMatches ? "" : searchValue} // When a parent directory matches the search, we don't need to filter the children.
          onOpenFile={onOpenFile}
        />
      )}
    </>
  );
};

const StorageNamespaceSection: React.FC<{
  namespace: StorageNamespace;
  locale: string;
  searchValue: string;
  onOpenFile: (info: OpenFileInfo) => void;
}> = ({ namespace, locale, searchValue, onOpenFile }) => {
  const [isExpanded, setIsExpanded] = useState(true);
  const { entriesByPath } = useStorage();
  const { clearNamespaceCache } = useStorageActions();
  const namespaceName = namespace.name ?? namespace.displayName;

  const {
    entries: fetchedEntries,
    isPending,
    error,
    refetch,
  } = useStorageEntries(namespaceName);

  const handleRefresh = useCallback(
    (e: React.MouseEvent) => {
      e.stopPropagation();
      clearNamespaceCache(namespaceName);
      refetch();
    },
    [namespaceName, clearNamespaceCache, refetch],
  );

  // While loading, fall back to initial entries from the namespace notification
  const entries = isPending ? namespace.storageEntries : fetchedEntries;
  const filtered = filterEntries(
    entries,
    namespaceName,
    searchValue,
    entriesByPath,
  );

  return (
    <>
      <CommandItem
        value={namespace.name}
        onSelect={() => setIsExpanded(!isExpanded)}
        className="flex flex-row font-semibold h-7 text-xs gap-1.5 bg-(--slate-2) text-muted-foreground rounded-none"
      >
        <TreeChevron isExpanded={isExpanded} className="h-3 w-3" />
        <ProtocolIcon protocol={namespace.protocol} />
        <span>{namespace.displayName}</span>
        {namespace.name && (
          <span className="text-xs text-muted-foreground font-normal">
            (<EngineVariable variableName={namespace.name} />)
          </span>
        )}
        <RefreshIconButton
          onClick={handleRefresh}
          tooltip="Refresh storage connection"
          className="p-0"
          iconClassName="h-3 w-3"
        />
        <span className="text-[10px] text-muted-foreground font-normal tabular-nums ml-auto">
          {namespace.rootPath || "(root)"}
        </span>
      </CommandItem>
      {isExpanded && (
        <>
          {isPending && entries.length === 0 && (
            <div
              className="flex items-center gap-1.5 py-1 text-xs text-muted-foreground"
              style={indentStyle(1)}
            >
              <LoaderCircle className="h-3 w-3 animate-spin" />
              Loading...
            </div>
          )}
          {error && entries.length === 0 && (
            <ErrorState
              error={error}
              style={indentStyle(1)}
              className="py-1 text-xs h-auto overflow-auto max-h-32 items-start"
              showIcon={false}
            />
          )}
          {!isPending && entries.length === 0 && !error && (
            <div
              className="py-1 text-xs text-muted-foreground italic"
              style={indentStyle(1)}
            >
              No entries
            </div>
          )}
          {searchValue && filtered.length === 0 && entries.length > 0 && (
            <div
              className="py-1 text-xs text-muted-foreground italic"
              style={indentStyle(1)}
            >
              No matches
            </div>
          )}
          {filtered.map((entry) => (
            <StorageEntryRow
              key={entry.path}
              entry={entry}
              namespace={namespaceName}
              protocol={namespace.protocol}
              rootPath={namespace.rootPath}
              backendType={namespace.backendType}
              depth={1}
              locale={locale}
              searchValue={searchValue}
              onOpenFile={onOpenFile}
            />
          ))}
        </>
      )}
    </>
  );
};

export const StorageInspector: React.FC = () => {
  const { namespaces } = useStorage();
  const { locale } = useLocale();
  const [searchValue, setSearchValue] = useState("");
  const [openFile, setOpenFile] = useState<OpenFileInfo | null>(null);
  const hasSearch = !!searchValue.trim();

  if (namespaces.length === 0) {
    return (
      <PanelEmptyState
        title="No storage connected"
        description={
          <span>
            Create an obstore or fsspec connection in your notebook. See the{" "}
            <a
              className="text-link"
              href="https://docs.marimo.io/guides/working_with_data/remote_storage/#quick-start"
              target="_blank"
              rel="noopener noreferrer"
            >
              docs
            </a>
            .
          </span>
        }
        action={
          <AddConnectionDialog defaultTab="storage">
            <Button variant="outline" size="sm">
              Add remote storage
              <PlusIcon className="h-4 w-4 ml-2" />
            </Button>
          </AddConnectionDialog>
        }
        icon={<HardDriveIcon className="h-8 w-8" />}
      />
    );
  }

  return (
    <div className="h-full flex flex-col">
      {openFile && (
        <StorageFileViewer
          entry={openFile.entry}
          namespace={openFile.namespace}
          protocol={openFile.protocol}
          backendType={openFile.backendType}
          onBack={() => setOpenFile(null)}
        />
      )}
      <Command
        className={cn(
          "border-b bg-background rounded-none h-full pb-10 overflow-auto outline-hidden scrollbar-thin",
          // We want to keep the command open but hidden when a file is opened to preserve state (folders opened etc.)
          // TODO: Preserve scroll position
          openFile && "hidden",
        )}
        shouldFilter={false}
      >
        <div className="flex items-center w-full border-b">
          <CommandInput
            placeholder="Search entries..."
            className="h-6 m-1"
            value={searchValue}
            onValueChange={setSearchValue}
            rootClassName="flex-1 border-b-0"
          />
          {hasSearch && (
            <Button
              variant="text"
              size="xs"
              className="float-right border-none px-2 m-0 h-full"
              onClick={() => setSearchValue("")}
            >
              <XIcon className="h-4 w-4" />
            </Button>
          )}
          <Tooltip
            content="Filters loaded entries only. Expand directories to include their contents in the search."
            delayDuration={200}
          >
            <HelpCircleIcon className="h-3.5 w-3.5 shrink-0 cursor-help text-muted-foreground hover:text-foreground mr-2" />
          </Tooltip>
          <AddConnectionDialog defaultTab="storage">
            <Button
              variant="ghost"
              size="sm"
              className="px-2 border-0 border-l border-muted-background rounded-none focus-visible:ring-0 focus-visible:ring-offset-0"
            >
              <PlusIcon className="h-4 w-4" />
            </Button>
          </AddConnectionDialog>
        </div>
        <CommandList className="flex flex-col">
          {namespaces.map((ns) => (
            <StorageNamespaceSection
              key={ns.name ?? ns.displayName}
              namespace={ns}
              locale={locale}
              searchValue={searchValue}
              onOpenFile={setOpenFile}
            />
          ))}
        </CommandList>
      </Command>
    </div>
  );
};