src/app/chat/_components/DesignIterationPanel.tsx

// DesignIterationPanel.tsx — Multi-file design iteration UI.
// Renders when an assistant message contains a CSV artifact with design-like
// columns (iPAE, pLDDT, affinity, consensus_score, etc.). Provides sortable
// columns, row selection, threshold filters, and re-submit actions that send
// a new user message into the chat.
"use client";

import { useCallback, useEffect, useMemo, useState } from "react";
import { parseCsv, detectNumericColumns, isDesignCsv, detectDelimiter } from "@/lib/csv-utils";
import type { ParsedCsv } from "@/lib/csv-utils";
import type { Artifact } from "@/lib/types";

// ---- Column scoring metadata ----

/** Columns where LOWER is better (e.g., iPAE < 8). */
const LOWER_IS_BETTER = new Set(["ipae", "interaction_pae"]);
/** Columns where HIGHER is better (e.g., pLDDT > 85). */
const HIGHER_IS_BETTER = new Set(["plddt", "pred_lddt", "affinity", "affinity_pm", "consensus_score"]);

type SortDir = "asc" | "desc";

interface Props {
  artifact: Artifact;
  /** Called with a message to send as the next user turn. */
  onSendMessage: (msg: string) => void;
}

export default function DesignIterationPanel({ artifact, onSendMessage }: Props) {
  const [raw, setRaw] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);

  // Sort state
  const [sortCol, setSortCol] = useState<number | null>(null);
  const [sortDir, setSortDir] = useState<SortDir>("asc");

  // Selection state
  const [selected, setSelected] = useState<Set<number>>(new Set());

  // Filter thresholds
  const [ipaeMax, setIpaeMax] = useState<number>(10);
  const [plddtMin, setPlddtMin] = useState<number>(80);
  const [filtersActive, setFiltersActive] = useState(false);

  // Top-N re-submit count
  const [topN, setTopN] = useState(5);

  const delim = useMemo(() => detectDelimiter(artifact.name, artifact.url), [artifact]);

  useEffect(() => {
    let cancelled = false;
    setRaw(null);
    setError(null);
    fetch(artifact.url)
      .then((r) => { if (!r.ok) throw new Error(`HTTP ${r.status}`); return r.text(); })
      .then((t) => { if (!cancelled) setRaw(t); })
      .catch((e) => { if (!cancelled) setError(e.message || String(e)); });
    return () => { cancelled = true; };
  }, [artifact.url]);

  const parsed = useMemo<ParsedCsv | null>(() => {
    if (raw == null) return null;
    try {
      const p = parseCsv(raw, delim);
      if (p.header.length === 0) return null;
      return p;
    } catch { return null; }
  }, [raw, delim]);

  const numericCols = useMemo(
    () => (parsed ? detectNumericColumns(parsed.header, parsed.rows) : []),
    [parsed],
  );

  // Column indices for design metrics
  const colMap = useMemo(() => {
    if (!parsed) return {} as Record<string, number>;
    const map: Record<string, number> = {};
    parsed.header.forEach((h, i) => {
      const key = h.toLowerCase().replace(/[^a-z0-9_]/g, "_");
      map[key] = i;
    });
    return map;
  }, [parsed]);

  const ipaeCol = colMap["ipae"] ?? colMap["interaction_pae"] ?? -1;
  const plddtCol = colMap["plddt"] ?? colMap["pred_lddt"] ?? -1;

  // Filtered rows (apply threshold filters when active)
  const filteredRows = useMemo(() => {
    if (!parsed) return [];
    let rows = parsed.rows.map((r, origIdx) => ({ row: r, origIdx }));

    if (filtersActive) {
      rows = rows.filter(({ row }) => {
        if (ipaeCol >= 0) {
          const v = parseFloat(row[ipaeCol] ?? "");
          if (!isNaN(v) && v > ipaeMax) return false;
        }
        if (plddtCol >= 0) {
          const v = parseFloat(row[plddtCol] ?? "");
          if (!isNaN(v) && v < plddtMin) return false;
        }
        return true;
      });
    }

    // Apply sort
    if (sortCol !== null && numericCols[sortCol]) {
      rows.sort((a, b) => {
        const va = parseFloat(a.row[sortCol] ?? "");
        const vb = parseFloat(b.row[sortCol] ?? "");
        if (isNaN(va) && isNaN(vb)) return 0;
        if (isNaN(va)) return 1;
        if (isNaN(vb)) return -1;
        return sortDir === "asc" ? va - vb : vb - va;
      });
    } else if (sortCol !== null) {
      rows.sort((a, b) => {
        const va = a.row[sortCol] ?? "";
        const vb = b.row[sortCol] ?? "";
        return sortDir === "asc" ? va.localeCompare(vb) : vb.localeCompare(va);
      });
    }

    return rows;
  }, [parsed, filtersActive, ipaeMax, plddtMin, ipaeCol, plddtCol, sortCol, sortDir, numericCols]);

  const handleSort = useCallback((colIdx: number) => {
    if (sortCol === colIdx) {
      setSortDir((d) => (d === "asc" ? "desc" : "asc"));
    } else {
      setSortCol(colIdx);
      // Default sort direction: lower-is-better → asc, higher-is-better → desc
      const key = parsed?.header[colIdx]?.toLowerCase().replace(/[^a-z0-9_]/g, "_") ?? "";
      if (LOWER_IS_BETTER.has(key)) setSortDir("asc");
      else if (HIGHER_IS_BETTER.has(key)) setSortDir("desc");
      else setSortDir("asc");
    }
  }, [sortCol, parsed]);

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

  const toggleSelectAll = useCallback(() => {
    if (selected.size === filteredRows.length) {
      setSelected(new Set());
    } else {
      setSelected(new Set(filteredRows.map((r) => r.origIdx)));
    }
  }, [selected, filteredRows]);

  // Build IDs from selected rows (use design_id/tag column or row index)
  const getDesignIds = useCallback((indices: number[]): string[] => {
    if (!parsed) return [];
    const idCol = colMap["design_id"] ?? colMap["tag"] ?? -1;
    return indices.map((idx) => {
      if (idCol >= 0) return parsed.rows[idx]?.[idCol] ?? `row_${idx}`;
      return `row_${idx}`;
    });
  }, [parsed, colMap]);

  const handleSubmitSelected = useCallback(() => {
    const ids = getDesignIds(Array.from(selected));
    if (ids.length === 0) return;
    onSendMessage(`Score these designs: ${ids.join(", ")}`);
  }, [selected, getDesignIds, onSendMessage]);

  const handleResubmitTop = useCallback(() => {
    const topRows = filteredRows.slice(0, topN);
    const ids = getDesignIds(topRows.map((r) => r.origIdx));
    if (ids.length === 0) return;

    const filterDesc: string[] = [];
    if (ipaeCol >= 0) filterDesc.push(`iPAE < ${ipaeMax}`);
    if (plddtCol >= 0) filterDesc.push(`pLDDT > ${plddtMin}`);
    const filterStr = filterDesc.length > 0 ? ` (filtered by ${filterDesc.join(" and ")})` : "";

    onSendMessage(
      `Filter ${artifact.name} by${filterStr} and re-run the top ${topN} designs with ProteinMPNN`,
    );
  }, [filteredRows, topN, getDesignIds, ipaeCol, plddtCol, ipaeMax, plddtMin, artifact.name, onSendMessage]);

  if (error) {
    return (
      <div className="p-3 text-xs text-red-600 border border-red-200 rounded-lg bg-red-50">
        Failed to load design CSV: {error}
      </div>
    );
  }

  if (raw == null) {
    return <div className="p-3 text-xs text-muted-fg">Loading design data...</div>;
  }

  if (!parsed || !isDesignCsv(parsed.header)) {
    return null; // Not a design CSV — don't render this panel
  }

  const allSelected = selected.size === filteredRows.length && filteredRows.length > 0;

  return (
    <div
      data-testid="design-iteration-panel"
      className="mt-3 rounded-xl border border-accent/20 bg-white shadow-sm overflow-hidden"
    >
      {/* Header */}
      <div className="flex items-center justify-between px-4 py-2.5 border-b border-border-subtle bg-accent/5">
        <div className="flex items-center gap-2">
          <svg className="w-4 h-4 text-accent" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth={2}>
            <path d="M2 15c6.667-6 13.333 0 20-6" />
            <path d="M2 9c6.667 6 13.333 0 20 6" />
          </svg>
          <span className="text-xs font-semibold text-foreground">Design Iteration</span>
          <span className="text-[10px] text-muted-fg">
            {filteredRows.length} design{filteredRows.length === 1 ? "" : "s"}
            {filtersActive && ` (filtered from ${parsed.rows.length})`}
          </span>
        </div>
        <div className="flex items-center gap-2">
          <button
            type="button"
            onClick={() => setFiltersActive((f) => !f)}
            className={`text-[10px] px-2 py-1 rounded-md transition-colors ${
              filtersActive
                ? "bg-accent text-white"
                : "text-muted-fg hover:text-foreground hover:bg-muted"
            }`}
          >
            <svg className="w-3 h-3 inline mr-1" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth={2}>
              <polygon points="22 3 2 3 10 12.46 10 19 14 21 14 12.46 22 3" />
            </svg>
            Filters
          </button>
        </div>
      </div>

      {/* Filter controls */}
      {filtersActive && (
        <div className="px-4 py-2.5 border-b border-border-subtle bg-muted/30 flex flex-wrap gap-4 items-center">
          {ipaeCol >= 0 && (
            <label className="flex items-center gap-2 text-[11px] text-foreground-dim">
              iPAE &lt;
              <input
                type="number"
                value={ipaeMax}
                onChange={(e) => setIpaeMax(parseFloat(e.target.value) || 0)}
                step={0.5}
                min={0}
                max={50}
                className="w-16 rounded-md border border-border bg-white px-2 py-1 text-xs focus:outline-none focus:border-accent"
              />
            </label>
          )}
          {plddtCol >= 0 && (
            <label className="flex items-center gap-2 text-[11px] text-foreground-dim">
              pLDDT &gt;
              <input
                type="number"
                value={plddtMin}
                onChange={(e) => setPlddtMin(parseFloat(e.target.value) || 0)}
                step={1}
                min={0}
                max={100}
                className="w-16 rounded-md border border-border bg-white px-2 py-1 text-xs focus:outline-none focus:border-accent"
              />
            </label>
          )}
          <div className="flex items-center gap-2 ml-auto text-[11px] text-foreground-dim">
            <span>Top</span>
            <input
              type="number"
              value={topN}
              onChange={(e) => setTopN(Math.max(1, parseInt(e.target.value) || 1))}
              min={1}
              max={filteredRows.length || 100}
              className="w-14 rounded-md border border-border bg-white px-2 py-1 text-xs focus:outline-none focus:border-accent"
            />
            <button
              type="button"
              onClick={handleResubmitTop}
              disabled={filteredRows.length === 0}
              className="px-2.5 py-1 rounded-md bg-accent text-white text-[10px] font-medium hover:bg-accent-hover disabled:opacity-30 transition-colors"
            >
              Re-submit top {Math.min(topN, filteredRows.length)}
            </button>
          </div>
        </div>
      )}

      {/* Table */}
      <div className="max-h-80 overflow-auto">
        <table className="w-full border-collapse font-mono text-[11px] leading-tight">
          <thead className="sticky top-0 z-10 bg-muted/80 backdrop-blur">
            <tr>
              <th className="px-2 py-1.5 text-left w-8">
                <input
                  type="checkbox"
                  checked={allSelected}
                  onChange={toggleSelectAll}
                  className="accent-accent"
                  aria-label="Select all designs"
                />
              </th>
              {parsed.header.map((h, idx) => {
                const isSorted = sortCol === idx;
                return (
                  <th
                    key={idx}
                    title={h}
                    onClick={() => handleSort(idx)}
                    className={`px-2 py-1.5 text-[10px] font-semibold uppercase tracking-wide border-b border-border-subtle whitespace-nowrap cursor-pointer select-none hover:text-accent transition-colors ${
                      numericCols[idx] ? "text-right" : "text-left"
                    } ${isSorted ? "text-accent" : "text-muted-fg"}`}
                    style={{ maxWidth: 140 }}
                  >
                    <span className="inline-flex items-center gap-1">
                      <span className="overflow-hidden text-ellipsis whitespace-nowrap">{h}</span>
                      {isSorted && (
                        <svg className={`w-3 h-3 shrink-0 ${sortDir === "desc" ? "rotate-180" : ""}`} viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth={2}>
                          <polyline points="18 15 12 9 6 15" />
                        </svg>
                      )}
                    </span>
                  </th>
                );
              })}
            </tr>
          </thead>
          <tbody>
            {filteredRows.map(({ row, origIdx }) => {
              const isSelected = selected.has(origIdx);
              return (
                <tr
                  key={origIdx}
                  className={`${isSelected ? "bg-accent/5" : origIdx % 2 === 1 ? "bg-muted/30" : ""} hover:bg-accent/10 transition-colors cursor-pointer`}
                  onClick={() => toggleSelect(origIdx)}
                >
                  <td className="px-2 py-1">
                    <input
                      type="checkbox"
                      checked={isSelected}
                      onChange={() => toggleSelect(origIdx)}
                      className="accent-accent"
                      aria-label={`Select design ${origIdx}`}
                      onClick={(e) => e.stopPropagation()}
                    />
                  </td>
                  {parsed.header.map((_, cIdx) => {
                    const val = row[cIdx] ?? "";
                    return (
                      <td
                        key={cIdx}
                        title={val}
                        className={`px-2 py-1 border-b border-border-subtle/60 ${
                          numericCols[cIdx] ? "text-right tabular-nums" : "text-left"
                        }`}
                        style={{ maxWidth: 140 }}
                      >
                        <span className="block overflow-hidden text-ellipsis whitespace-nowrap">
                          {val}
                        </span>
                      </td>
                    );
                  })}
                </tr>
              );
            })}
          </tbody>
        </table>
      </div>

      {/* Action footer */}
      <div className="flex items-center justify-between px-4 py-2.5 border-t border-border-subtle bg-muted/20">
        <span className="text-[10px] text-muted-fg">
          {selected.size} of {filteredRows.length} selected
        </span>
        <button
          type="button"
          onClick={handleSubmitSelected}
          disabled={selected.size === 0}
          className="px-3 py-1.5 rounded-md bg-accent text-white text-[10px] font-medium hover:bg-accent-hover disabled:opacity-30 transition-colors"
        >
          Submit selected for scoring
        </button>
      </div>
    </div>
  );
}

export { DesignIterationPanel };