frontend/src/adaevolve/AdaevolveApp.tsx

import { useEffect, useState, useMemo } from "react";
import { useAppState } from "./store";
import type { IterationSummary } from "./types";

const ADAPTATION_COLORS: Record<string, { bg: string; text: string; bar: string }> = {
  L1_explore: { bg: "bg-blue-900/40", text: "text-blue-400", bar: "bg-blue-500" },
  L1_exploit: { bg: "bg-green-900/40", text: "text-green-400", bar: "bg-green-500" },
  L2_migrate: { bg: "bg-amber-900/40", text: "text-amber-400", bar: "bg-amber-500" },
  L3_meta: { bg: "bg-red-900/40", text: "text-red-400", bar: "bg-red-500" },
};

const DEFAULT_COLOR = { bg: "bg-gray-900/40", text: "text-gray-400", bar: "bg-gray-500" };

function getAdaptationColor(type: string) {
  return ADAPTATION_COLORS[type] || DEFAULT_COLOR;
}

type SortKey = "iteration" | "score" | "best_score" | "delta" | "island_id" | "adaptation_type";
type SortDir = "asc" | "desc";

export default function AdaevolveApp() {
  const store = useAppState();
  const { state } = store;
  const [repoInput, setRepoInput] = useState("");
  const [sortKey, setSortKey] = useState<SortKey>("iteration");
  const [sortDir, setSortDir] = useState<SortDir>("asc");

  useEffect(() => {
    store.loadPresets();
  }, []);

  // Get the currently selected dataset
  const selectedDataset = useMemo(
    () => state.datasets.find((d) => d.id === state.selectedDatasetId) ?? null,
    [state.datasets, state.selectedDatasetId]
  );

  // Filter iterations by adaptation type
  const filteredIterations = useMemo(() => {
    if (!selectedDataset) return [];
    let iters = selectedDataset.iterations;
    if (state.filterAdaptation) {
      iters = iters.filter((it) => it.adaptation_type === state.filterAdaptation);
    }
    return iters;
  }, [selectedDataset, state.filterAdaptation]);

  // Sort iterations
  const sortedIterations = useMemo(() => {
    const sorted = [...filteredIterations];
    sorted.sort((a, b) => {
      const av = a[sortKey];
      const bv = b[sortKey];
      if (typeof av === "number" && typeof bv === "number") {
        return sortDir === "asc" ? av - bv : bv - av;
      }
      const sa = String(av);
      const sb = String(bv);
      return sortDir === "asc" ? sa.localeCompare(sb) : sb.localeCompare(sa);
    });
    return sorted;
  }, [filteredIterations, sortKey, sortDir]);

  // Get all unique adaptation types across selected dataset
  const adaptationTypes = useMemo(() => {
    if (!selectedDataset) return [];
    const types = new Set(selectedDataset.iterations.map((it) => it.adaptation_type));
    return Array.from(types).sort();
  }, [selectedDataset]);

  function handleSort(key: SortKey) {
    if (sortKey === key) {
      setSortDir((d) => (d === "asc" ? "desc" : "asc"));
    } else {
      setSortKey(key);
      setSortDir("asc");
    }
  }

  function handleLoad() {
    if (repoInput.trim()) {
      store.loadDataset(repoInput.trim());
      setRepoInput("");
    }
  }

  // Compute max score for bar chart scaling
  const maxScore = useMemo(() => {
    if (!filteredIterations.length) return 1;
    return Math.max(...filteredIterations.map((it) => Math.abs(it.score)), 0.001);
  }, [filteredIterations]);

  const sortIndicator = (key: SortKey) => {
    if (sortKey !== key) return "";
    return sortDir === "asc" ? " ^" : " v";
  };

  return (
    <div className="flex h-full overflow-hidden">
      {/* Sidebar */}
      <div className="w-72 border-r border-gray-800 bg-gray-900 flex flex-col overflow-hidden shrink-0">
        {/* Repo input */}
        <div className="p-3 border-b border-gray-800">
          <label className="text-xs text-gray-500 mb-1 block">HuggingFace Repo</label>
          <div className="flex gap-1">
            <input
              type="text"
              value={repoInput}
              onChange={(e) => setRepoInput(e.target.value)}
              onKeyDown={(e) => e.key === "Enter" && handleLoad()}
              placeholder="org/dataset-name"
              className="flex-1 bg-gray-800 text-sm px-2 py-1.5 rounded border border-gray-700 focus:border-rose-500 focus:outline-none text-gray-200"
            />
            <button
              onClick={handleLoad}
              disabled={state.loading || !repoInput.trim()}
              className="bg-rose-600 hover:bg-rose-500 disabled:opacity-50 text-white text-sm px-3 py-1.5 rounded"
            >
              Load
            </button>
          </div>
        </div>

        {/* Presets */}
        <div className="p-3 border-b border-gray-800">
          <div className="text-xs text-gray-500 mb-2">Presets</div>
          {state.presets.length === 0 && (
            <div className="text-xs text-gray-600">No presets available</div>
          )}
          <div className="space-y-1 max-h-32 overflow-y-auto">
            {state.presets.map((p) => (
              <div
                key={p.id}
                className="flex items-center justify-between group"
              >
                <button
                  onClick={() => store.loadPreset(p)}
                  className="text-xs text-rose-400 hover:text-rose-300 truncate flex-1 text-left"
                  title={p.repo}
                >
                  {p.name}
                </button>
                <button
                  onClick={() => store.deletePreset(p.id)}
                  className="text-gray-600 hover:text-red-400 text-xs opacity-0 group-hover:opacity-100 ml-1"
                >
                  x
                </button>
              </div>
            ))}
          </div>
        </div>

        {/* Loaded datasets */}
        <div className="flex-1 overflow-y-auto p-3">
          <div className="text-xs text-gray-500 mb-2">
            Loaded Datasets ({state.datasets.length})
          </div>
          {state.datasets.map((ds) => (
            <div
              key={ds.id}
              className={`mb-3 p-2 rounded border cursor-pointer ${
                state.selectedDatasetId === ds.id
                  ? "border-rose-500 bg-rose-900/20"
                  : "border-gray-700 hover:border-gray-600"
              }`}
              onClick={() => store.selectDataset(ds.id)}
            >
              <div className="flex items-center justify-between">
                <span className="text-sm text-gray-200 truncate" title={ds.repo}>
                  {ds.name}
                </span>
                <button
                  onClick={(e) => {
                    e.stopPropagation();
                    store.unloadDataset(ds.id);
                  }}
                  className="text-gray-600 hover:text-red-400 text-xs ml-1"
                >
                  x
                </button>
              </div>
              <div className="text-xs text-gray-500 mt-1">
                {ds.n_iterations} iterations | {ds.summary_stats.n_islands} islands
              </div>
              <div className="text-xs text-gray-500">
                Best: {ds.summary_stats.global_best.toFixed(4)}
              </div>
              {/* Adaptation breakdown */}
              <div className="mt-1 space-y-0.5">
                {Object.entries(ds.summary_stats.adaptation_counts).map(([type, count]) => {
                  const color = getAdaptationColor(type);
                  return (
                    <div key={type} className="flex items-center text-xs gap-1">
                      <span className={`w-2 h-2 rounded-full ${color.bar}`} />
                      <span className={color.text}>{type}</span>
                      <span className="text-gray-600 ml-auto">{count}</span>
                    </div>
                  );
                })}
              </div>
            </div>
          ))}
        </div>

        {/* Filter by adaptation type */}
        {selectedDataset && (
          <div className="p-3 border-t border-gray-800">
            <div className="text-xs text-gray-500 mb-1">Filter by Adaptation</div>
            <select
              value={state.filterAdaptation}
              onChange={(e) => store.setFilterAdaptation(e.target.value)}
              className="w-full bg-gray-800 text-sm px-2 py-1 rounded border border-gray-700 text-gray-200"
            >
              <option value="">All</option>
              {adaptationTypes.map((t) => (
                <option key={t} value={t}>
                  {t}
                </option>
              ))}
            </select>
          </div>
        )}
      </div>

      {/* Main content */}
      <div className="flex-1 flex flex-col overflow-hidden">
        {state.error && (
          <div className="bg-red-900/50 border-b border-red-700 px-4 py-2 text-sm text-red-200 flex justify-between items-center">
            <span>{state.error}</span>
            <button
              onClick={() => store.setError(null)}
              className="text-red-400 hover:text-red-200 ml-4"
            >
              dismiss
            </button>
          </div>
        )}

        {state.loading && (
          <div className="bg-rose-900/30 border-b border-rose-800 px-4 py-1.5 text-xs text-rose-300">
            Loading...
          </div>
        )}

        {/* Timeline view */}
        {state.viewMode === "timeline" && selectedDataset && (
          <div className="flex-1 flex flex-col overflow-hidden">
            {/* Bar chart */}
            <div className="border-b border-gray-800 p-4 shrink-0">
              <div className="text-sm text-gray-400 mb-2">
                Score Timeline ({filteredIterations.length} iterations)
              </div>
              <div className="flex items-end gap-px h-32 overflow-x-auto">
                {filteredIterations.map((it) => {
                  const height = Math.max((Math.abs(it.score) / maxScore) * 100, 2);
                  const color = getAdaptationColor(it.adaptation_type);
                  return (
                    <div
                      key={it.index}
                      className="flex flex-col items-center justify-end flex-shrink-0 cursor-pointer group"
                      style={{ minWidth: Math.max(4, Math.min(20, 800 / filteredIterations.length)) }}
                      onClick={() => store.selectIteration(selectedDataset.id, it.index)}
                      title={`iter ${it.iteration} | score ${it.score.toFixed(4)} | ${it.adaptation_type}`}
                    >
                      <div
                        className={`w-full rounded-t ${color.bar} group-hover:opacity-80 transition-opacity`}
                        style={{ height: `${height}%` }}
                      />
                    </div>
                  );
                })}
              </div>
              {/* Legend */}
              <div className="flex gap-4 mt-2">
                {Object.entries(ADAPTATION_COLORS).map(([type, color]) => (
                  <div key={type} className="flex items-center gap-1 text-xs">
                    <span className={`w-2 h-2 rounded-full ${color.bar}`} />
                    <span className={color.text}>{type}</span>
                  </div>
                ))}
              </div>
            </div>

            {/* Iteration table */}
            <div className="flex-1 overflow-auto">
              <table className="w-full text-sm">
                <thead className="sticky top-0 bg-gray-900 border-b border-gray-800">
                  <tr>
                    {(
                      [
                        ["iteration", "Iter"],
                        ["island_id", "Island"],
                        ["score", "Score"],
                        ["best_score", "Best"],
                        ["delta", "Delta"],
                        ["adaptation_type", "Adaptation"],
                      ] as [SortKey, string][]
                    ).map(([key, label]) => (
                      <th
                        key={key}
                        onClick={() => handleSort(key)}
                        className="px-3 py-2 text-left text-gray-400 font-medium cursor-pointer hover:text-gray-200 select-none"
                      >
                        {label}
                        {sortIndicator(key)}
                      </th>
                    ))}
                    <th className="px-3 py-2 text-left text-gray-400 font-medium">Valid</th>
                    <th className="px-3 py-2 text-left text-gray-400 font-medium">Task</th>
                  </tr>
                </thead>
                <tbody>
                  {sortedIterations.map((it) => {
                    const color = getAdaptationColor(it.adaptation_type);
                    return (
                      <tr
                        key={it.index}
                        onClick={() => store.selectIteration(selectedDataset.id, it.index)}
                        className="border-b border-gray-800/50 hover:bg-gray-800/50 cursor-pointer"
                      >
                        <td className="px-3 py-1.5 text-gray-300">{it.iteration}</td>
                        <td className="px-3 py-1.5 text-gray-300">{it.island_id}</td>
                        <td className="px-3 py-1.5 text-gray-200 font-mono">
                          {it.score.toFixed(4)}
                        </td>
                        <td className="px-3 py-1.5 text-gray-200 font-mono">
                          {it.best_score.toFixed(4)}
                        </td>
                        <td
                          className={`px-3 py-1.5 font-mono ${
                            it.delta > 0
                              ? "text-green-400"
                              : it.delta < 0
                              ? "text-red-400"
                              : "text-gray-500"
                          }`}
                        >
                          {it.delta > 0 ? "+" : ""}
                          {it.delta.toFixed(4)}
                        </td>
                        <td className="px-3 py-1.5">
                          <span
                            className={`px-1.5 py-0.5 rounded text-xs ${color.bg} ${color.text}`}
                          >
                            {it.adaptation_type}
                          </span>
                        </td>
                        <td className="px-3 py-1.5">
                          {it.is_valid ? (
                            <span className="text-green-400 text-xs">yes</span>
                          ) : (
                            <span className="text-red-400 text-xs">no</span>
                          )}
                        </td>
                        <td className="px-3 py-1.5 text-gray-500 text-xs truncate max-w-[150px]">
                          {it.task_id}
                        </td>
                      </tr>
                    );
                  })}
                </tbody>
              </table>
            </div>
          </div>
        )}

        {/* Detail view */}
        {state.viewMode === "detail" && state.iterationDetail && (
          <div className="flex-1 flex flex-col overflow-hidden">
            {/* Detail header */}
            <div className="border-b border-gray-800 px-4 py-2 flex items-center gap-4 shrink-0">
              <button
                onClick={() => store.backToTimeline()}
                className="text-rose-400 hover:text-rose-300 text-sm"
              >
                &larr; Back to timeline
              </button>
              <div className="text-sm text-gray-300">
                Iteration {state.iterationDetail.iteration} | Island{" "}
                {state.iterationDetail.island_id}
              </div>
              <div className="flex items-center gap-2 ml-auto">
                <span
                  className={`px-2 py-0.5 rounded text-xs ${
                    getAdaptationColor(state.iterationDetail.adaptation_type).bg
                  } ${getAdaptationColor(state.iterationDetail.adaptation_type).text}`}
                >
                  {state.iterationDetail.adaptation_type}
                </span>
                <span className="text-gray-400 text-sm font-mono">
                  score: {state.iterationDetail.score.toFixed(4)}
                </span>
                <span
                  className={`text-sm font-mono ${
                    state.iterationDetail.delta > 0
                      ? "text-green-400"
                      : state.iterationDetail.delta < 0
                      ? "text-red-400"
                      : "text-gray-500"
                  }`}
                >
                  ({state.iterationDetail.delta > 0 ? "+" : ""}
                  {state.iterationDetail.delta.toFixed(4)})
                </span>
                {state.iterationDetail.is_valid ? (
                  <span className="text-green-400 text-xs border border-green-800 rounded px-1">
                    valid
                  </span>
                ) : (
                  <span className="text-red-400 text-xs border border-red-800 rounded px-1">
                    invalid
                  </span>
                )}
              </div>
            </div>

            {/* Detail grid */}
            <div className="flex-1 overflow-auto p-4">
              <div className="grid grid-cols-2 gap-4 h-full">
                {/* Meta info */}
                <div className="col-span-2 grid grid-cols-4 gap-3">
                  <div className="bg-gray-900 rounded border border-gray-800 p-3">
                    <div className="text-xs text-gray-500 mb-1">Task ID</div>
                    <div className="text-sm text-gray-200 break-all">
                      {state.iterationDetail.task_id || "-"}
                    </div>
                  </div>
                  <div className="bg-gray-900 rounded border border-gray-800 p-3">
                    <div className="text-xs text-gray-500 mb-1">Exploration Intensity</div>
                    <div className="text-sm text-gray-200 font-mono">
                      {state.iterationDetail.exploration_intensity.toFixed(4)}
                    </div>
                  </div>
                  <div className="bg-gray-900 rounded border border-gray-800 p-3">
                    <div className="text-xs text-gray-500 mb-1">Meta-Guidance Tactic</div>
                    <div className="text-sm text-gray-200">
                      {state.iterationDetail.meta_guidance_tactic || "-"}
                    </div>
                  </div>
                  <div className="bg-gray-900 rounded border border-gray-800 p-3">
                    <div className="text-xs text-gray-500 mb-1">Tactic Approach</div>
                    <div className="text-sm text-gray-200">
                      {state.iterationDetail.tactic_approach_type || "-"}
                    </div>
                  </div>
                </div>

                {/* Prompt */}
                <div className="bg-gray-900 rounded border border-gray-800 flex flex-col overflow-hidden">
                  <div className="px-3 py-2 border-b border-gray-800 text-xs text-rose-400 font-medium shrink-0">
                    Prompt
                  </div>
                  <pre className="flex-1 overflow-auto p-3 text-xs text-gray-300 whitespace-pre-wrap font-mono">
                    {state.iterationDetail.prompt_text || "(empty)"}
                  </pre>
                </div>

                {/* Reasoning trace */}
                <div className="bg-gray-900 rounded border border-gray-800 flex flex-col overflow-hidden">
                  <div className="px-3 py-2 border-b border-gray-800 text-xs text-rose-400 font-medium shrink-0">
                    Reasoning Trace
                  </div>
                  <pre className="flex-1 overflow-auto p-3 text-xs text-gray-300 whitespace-pre-wrap font-mono">
                    {state.iterationDetail.reasoning_trace || "(empty)"}
                  </pre>
                </div>

                {/* Generated code */}
                <div className="col-span-2 bg-gray-900 rounded border border-gray-800 flex flex-col overflow-hidden max-h-96">
                  <div className="px-3 py-2 border-b border-gray-800 text-xs text-rose-400 font-medium shrink-0">
                    Generated Code
                  </div>
                  <pre className="flex-1 overflow-auto p-3 text-xs text-gray-300 whitespace-pre-wrap font-mono">
                    {state.iterationDetail.program_code || "(empty)"}
                  </pre>
                </div>
              </div>
            </div>

            {/* Navigation */}
            {selectedDataset && (
              <div className="border-t border-gray-800 px-4 py-2 flex items-center justify-between shrink-0">
                <button
                  onClick={() => {
                    const idx = state.iterationDetail!.index;
                    if (idx > 0) store.selectIteration(selectedDataset.id, idx - 1);
                  }}
                  disabled={state.iterationDetail.index <= 0}
                  className="text-sm text-rose-400 hover:text-rose-300 disabled:opacity-30 disabled:cursor-default"
                >
                  &larr; Previous
                </button>
                <span className="text-xs text-gray-500">
                  {state.iterationDetail.index + 1} / {selectedDataset.n_iterations}
                </span>
                <button
                  onClick={() => {
                    const idx = state.iterationDetail!.index;
                    if (idx < selectedDataset.n_iterations - 1)
                      store.selectIteration(selectedDataset.id, idx + 1);
                  }}
                  disabled={
                    state.iterationDetail.index >= selectedDataset.n_iterations - 1
                  }
                  className="text-sm text-rose-400 hover:text-rose-300 disabled:opacity-30 disabled:cursor-default"
                >
                  Next &rarr;
                </button>
              </div>
            )}
          </div>
        )}

        {/* Empty state */}
        {state.datasets.length === 0 && state.viewMode === "timeline" && (
          <div className="flex-1 flex items-center justify-center text-gray-500">
            <div className="text-center">
              <div className="text-lg mb-2">No datasets loaded</div>
              <div className="text-sm">
                Load a HuggingFace repo from the sidebar or select a preset
              </div>
            </div>
          </div>
        )}

        {state.datasets.length > 0 && !selectedDataset && state.viewMode === "timeline" && (
          <div className="flex-1 flex items-center justify-center text-gray-500">
            <div className="text-center">
              <div className="text-lg mb-2">Select a dataset</div>
              <div className="text-sm">Click a dataset in the sidebar to view its iterations</div>
            </div>
          </div>
        )}
      </div>
    </div>
  );
}