src/app/chat/_components/DatasetSearchPanel.tsx

"use client";

import { useCallback, useMemo, useState } from "react";
import {
  searchDatasetsAPI,
  searchDatasetSourceAPI,
  importDatasetEntriesAPI,
} from "@/lib/api";
import type { Artifact, DatasetSearchResult } from "@/lib/types";

type SourceFilter = "all" | "sabdab" | "oas" | "imgt";

const SOURCE_LABELS: Record<string, string> = {
  sabdab: "SAbDab",
  oas: "OAS",
  imgt: "IMGT",
  fc: "Fc",
};

const DATA_TYPE_COLORS: Record<string, string> = {
  structure: "bg-blue-100 text-blue-700",
  sequence: "bg-emerald-100 text-emerald-700",
  gene: "bg-purple-100 text-purple-700",
};

interface Props {
  open: boolean;
  onClose: () => void;
  onImport: (artifacts: Artifact[]) => void;
}

export default function DatasetSearchPanel({ open, onClose, onImport }: Props) {
  const [query, setQuery] = useState("");
  const [source, setSource] = useState<SourceFilter>("all");
  const [results, setResults] = useState<Record<string, DatasetSearchResult[]>>(
    {},
  );
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [selected, setSelected] = useState<Set<string>>(new Set());
  const [importing, setImporting] = useState(false);
  const [collapsed, setCollapsed] = useState<Set<string>>(new Set());

  const search = useCallback(async () => {
    if (!query.trim()) return;
    setLoading(true);
    setError(null);
    setResults({});
    setSelected(new Set());
    try {
      if (source === "all") {
        const res = await searchDatasetsAPI({ query: query.trim() });
        setResults(res.results);
      } else {
        const res = await searchDatasetSourceAPI(source, {
          query: query.trim(),
        });
        setResults({ [source]: res });
      }
    } catch (e) {
      setError(e instanceof Error ? e.message : "Search failed");
    } finally {
      setLoading(false);
    }
  }, [query, source]);

  const handleKeyDown = useCallback(
    (e: React.KeyboardEvent) => {
      if (e.key === "Enter") {
        e.preventDefault();
        search();
      }
    },
    [search],
  );

  const toggleSelect = useCallback((id: string) => {
    setSelected((prev) => {
      const next = new Set(prev);
      if (next.has(id)) next.delete(id);
      else next.add(id);
      return next;
    });
  }, []);

  const toggleCollapse = useCallback((src: string) => {
    setCollapsed((prev) => {
      const next = new Set(prev);
      if (next.has(src)) next.delete(src);
      else next.add(src);
      return next;
    });
  }, []);

  const allResults = useMemo(() => {
    const flat: DatasetSearchResult[] = [];
    for (const items of Object.values(results)) {
      flat.push(...items);
    }
    return flat;
  }, [results]);

  const handleImport = useCallback(async () => {
    const entries = allResults
      .filter((r) => selected.has(r.id))
      .map((r) => ({ source: r.source, entry_id: r.id }));
    if (entries.length === 0) return;
    setImporting(true);
    try {
      const res = await importDatasetEntriesAPI(entries);
      onImport(res.artifacts);
    } catch (e) {
      setError(e instanceof Error ? e.message : "Import failed");
    } finally {
      setImporting(false);
    }
  }, [allResults, selected, onImport]);

  if (!open) return null;

  const sourceKeys = Object.keys(results).filter(
    (k) => results[k].length > 0,
  );
  const hasResults = sourceKeys.length > 0;

  return (
    <div
      role="dialog"
      aria-modal="true"
      aria-label="Search datasets"
      data-testid="dataset-search-panel"
      className="fixed inset-0 z-50 flex items-center justify-center bg-black/30 backdrop-blur-sm"
      onClick={onClose}
    >
      <div
        className="w-[640px] max-w-[94vw] max-h-[85vh] bg-white rounded-xl border border-border shadow-xl flex flex-col"
        onClick={(e) => e.stopPropagation()}
      >
        {/* Header */}
        <div className="px-4 py-3 border-b border-border flex items-center gap-3">
          <div>
            <p className="text-xs font-semibold text-foreground">
              Search Databases
            </p>
            <p className="text-[10px] text-muted-fg">
              Query SAbDab, OAS, and IMGT for antibody data
            </p>
          </div>
          <div className="flex-1" />
          <button
            onClick={onClose}
            className="text-[11px] text-muted-fg hover:text-foreground"
            aria-label="Close"
          >
            ✕
          </button>
        </div>

        {/* Search bar */}
        <div className="px-4 py-2 border-b border-border flex gap-2">
          <input
            type="text"
            value={query}
            onChange={(e) => setQuery(e.target.value)}
            onKeyDown={handleKeyDown}
            placeholder="Search antibody databases..."
            data-testid="dataset-search-input"
            className="flex-1 rounded-md border border-border bg-white px-3 py-1.5 text-xs focus:outline-none focus:border-accent focus:ring-1 focus:ring-accent/20"
          />
          <select
            value={source}
            onChange={(e) => setSource(e.target.value as SourceFilter)}
            data-testid="dataset-source-filter"
            className="rounded-md border border-border bg-white px-2 py-1.5 text-xs focus:outline-none focus:border-accent"
          >
            <option value="all">All Sources</option>
            <option value="sabdab">SAbDab</option>
            <option value="oas">OAS</option>
            <option value="imgt">IMGT</option>
          </select>
          <button
            onClick={search}
            disabled={loading || !query.trim()}
            className="px-3 py-1.5 rounded-md bg-accent text-white text-xs font-medium hover:bg-accent-hover disabled:opacity-30 transition-colors"
          >
            Search
          </button>
        </div>

        {/* Results */}
        <div className="flex-1 overflow-y-auto px-2 py-2">
          {loading && (
            <p className="text-[11px] text-muted-fg px-2 py-6 text-center">
              Searching...
            </p>
          )}
          {error && (
            <p className="text-[11px] text-red-600 px-2 py-6 text-center">
              {error}
            </p>
          )}
          {!loading && !error && !hasResults && query && (
            <p className="text-[11px] text-muted-fg px-2 py-6 text-center">
              No results found.
            </p>
          )}
          {!loading && !error && !query && !hasResults && (
            <p className="text-[11px] text-muted-fg px-2 py-6 text-center">
              Enter a query to search antibody databases.
            </p>
          )}

          {sourceKeys.map((src) => (
            <div key={src} className="mb-2">
              <button
                onClick={() => toggleCollapse(src)}
                className="w-full flex items-center gap-2 px-3 py-1.5 text-left"
              >
                <svg
                  className={`w-3 h-3 text-muted-fg transition-transform ${
                    collapsed.has(src) ? "" : "rotate-90"
                  }`}
                  viewBox="0 0 24 24"
                  fill="none"
                  stroke="currentColor"
                  strokeWidth={2}
                >
                  <polyline points="9 18 15 12 9 6" />
                </svg>
                <span className="text-[10px] font-bold uppercase tracking-widest text-muted-fg">
                  {SOURCE_LABELS[src] ?? src}
                </span>
                <span className="text-[10px] text-muted-fg">
                  ({results[src].length})
                </span>
              </button>

              {!collapsed.has(src) && (
                <ul className="space-y-0.5 px-1">
                  {results[src].map((r) => (
                    <li
                      key={r.id}
                      data-testid="dataset-result-row"
                      className="rounded-lg border border-border-subtle bg-white hover:border-accent/30 transition-colors"
                    >
                      <div className="flex items-start gap-2 px-3 py-2">
                        <input
                          type="checkbox"
                          checked={selected.has(r.id)}
                          onChange={() => toggleSelect(r.id)}
                          className="mt-0.5 shrink-0 accent-accent"
                          aria-label={`Select ${r.title}`}
                        />
                        <div className="flex-1 min-w-0">
                          <div className="flex items-baseline gap-2">
                            <span className="text-xs font-medium text-foreground truncate">
                              {r.title}
                            </span>
                            <span
                              className={`text-[9px] px-1.5 py-0.5 rounded-full font-medium ${
                                DATA_TYPE_COLORS[r.data_type] ??
                                "bg-gray-100 text-gray-600"
                              }`}
                            >
                              {r.data_type}
                            </span>
                          </div>
                          <p className="text-[11px] text-foreground-dim mt-0.5 leading-relaxed line-clamp-2">
                            {r.summary}
                          </p>
                          {r.organism && (
                            <p className="text-[10px] text-muted-fg mt-0.5">
                              {r.organism}
                            </p>
                          )}
                        </div>
                      </div>
                    </li>
                  ))}
                </ul>
              )}
            </div>
          ))}
        </div>

        {/* Footer with import */}
        {hasResults && (
          <div className="px-4 py-2 border-t border-border flex items-center justify-between">
            <p className="text-[10px] text-muted-fg">
              {selected.size} of {allResults.length} selected
            </p>
            <button
              onClick={handleImport}
              disabled={selected.size === 0 || importing}
              data-testid="dataset-import-button"
              className="px-3 py-1.5 rounded-md bg-accent text-white text-xs font-medium hover:bg-accent-hover disabled:opacity-30 transition-colors"
            >
              {importing ? "Importing..." : "Import Selected"}
            </button>
          </div>
        )}
      </div>
    </div>
  );
}