frontend/src/rlm-eval/store.ts

import { useState, useCallback, useEffect, useMemo } from "react";
import type {
  DatasetInfo,
  Preset,
  PanelNav,
  OverviewData,
  ExampleDetailData,
  RlmIterDetail,
} from "./types";
import { api } from "./api";

export function useAppState() {
  const [datasets, setDatasets] = useState<DatasetInfo[]>([]);
  const [presets, setPresets] = useState<Preset[]>([]);
  const [error, setError] = useState<string | null>(null);
  const [loading, setLoading] = useState<Record<string, boolean>>({});

  // Dual panel navigation
  const [panelA, setPanelA] = useState<PanelNav | null>(null);
  const [panelB, setPanelB] = useState<PanelNav | null>(null);
  const [comparisonMode, setComparisonMode] = useState(false);

  // Data caches
  const [overviewCache, setOverviewCache] = useState<Record<string, OverviewData>>({});
  const [exampleDetailCache, setExampleDetailCache] = useState<Record<string, ExampleDetailData>>({});
  const [iterDetailCache, setIterDetailCache] = useState<Record<string, RlmIterDetail>>({});

  // Load presets on mount
  useEffect(() => {
    api.listPresets().then((data) => setPresets(data as unknown as Preset[])).catch(() => {});
  }, []);

  const activeDatasets = useMemo(() => datasets.filter((d) => d.active), [datasets]);

  // Data fetching helpers
  const fetchOverview = useCallback(async (dsId: string) => {
    if (overviewCache[dsId]) return overviewCache[dsId];
    const data = (await api.getOverview(dsId)) as unknown as OverviewData;
    setOverviewCache((prev) => ({ ...prev, [dsId]: data }));
    return data;
  }, [overviewCache]);

  const fetchExampleDetail = useCallback(
    async (dsId: string, exampleIdx: number) => {
      const key = `${dsId}:${exampleIdx}`;
      if (exampleDetailCache[key]) return exampleDetailCache[key];
      const data = (await api.getExampleDetail(dsId, exampleIdx)) as unknown as ExampleDetailData;
      setExampleDetailCache((prev) => ({ ...prev, [key]: data }));
      return data;
    },
    [exampleDetailCache]
  );

  const fetchIterDetail = useCallback(
    async (dsId: string, exampleIdx: number, rlmIter: number) => {
      const key = `${dsId}:${exampleIdx}:${rlmIter}`;
      if (iterDetailCache[key]) return iterDetailCache[key];
      const data = (await api.getIterDetail(dsId, exampleIdx, rlmIter)) as unknown as RlmIterDetail;
      setIterDetailCache((prev) => ({ ...prev, [key]: data }));
      return data;
    },
    [iterDetailCache]
  );

  // Dataset operations
  const addDataset = useCallback(
    async (repo: string, config?: string, split?: string, presetId?: string, presetName?: string) => {
      setLoading((prev) => ({ ...prev, [repo]: true }));
      setError(null);
      try {
        const result = await api.loadDataset(repo, config, split);
        const dsInfo: DatasetInfo = {
          id: result.id,
          repo: result.repo,
          name: result.name,
          config: result.config,
          split: result.split,
          metadata: result.metadata as unknown as DatasetInfo["metadata"],
          n_examples: result.n_examples,
          n_rows: result.n_rows,
          active: true,
          presetId,
          presetName,
        };

        setDatasets((prev) => {
          if (prev.some((d) => d.id === dsInfo.id)) return prev;
          return [...prev, dsInfo];
        });

        // Auto-set panel A if not set
        setPanelA((prev) => prev || { datasetId: dsInfo.id, level: 1 });
      } catch (e: unknown) {
        setError(e instanceof Error ? e.message : "Failed to load dataset");
      } finally {
        setLoading((prev) => ({ ...prev, [repo]: false }));
      }
    },
    []
  );

  const removeDataset = useCallback(async (id: string) => {
    await api.unloadDataset(id).catch(() => {});
    setDatasets((prev) => prev.filter((d) => d.id !== id));
    setPanelA((prev) => (prev?.datasetId === id ? null : prev));
    setPanelB((prev) => (prev?.datasetId === id ? null : prev));
  }, []);

  const toggleDataset = useCallback((id: string) => {
    setDatasets((prev) => prev.map((d) => (d.id === id ? { ...d, active: !d.active } : d)));
  }, []);

  // Navigation
  const navigatePanel = useCallback(
    (panel: "A" | "B", nav: PanelNav) => {
      if (panel === "A") setPanelA(nav);
      else setPanelB(nav);
    },
    []
  );

  const goUp = useCallback((panel: "A" | "B") => {
    const setter = panel === "A" ? setPanelA : setPanelB;
    setter((prev) => {
      if (!prev) return prev;
      if (prev.level === 2) return { ...prev, level: 1, exampleIdx: undefined };
      return prev;
    });
  }, []);

  const updateDatasetPresetName = useCallback((dsId: string, name: string) => {
    setDatasets((prev) => prev.map((d) => (d.id === dsId ? { ...d, presetName: name } : d)));
  }, []);

  const clearDatasetPreset = useCallback((dsId: string) => {
    setDatasets((prev) =>
      prev.map((d) => (d.id === dsId ? { ...d, presetId: undefined, presetName: undefined } : d))
    );
  }, []);

  const toggleComparison = useCallback(() => {
    setComparisonMode((prev) => {
      if (!prev && panelA) {
        // Entering comparison: initialize panel B same as A
        setPanelB({ ...panelA });
      } else if (prev) {
        setPanelB(null);
      }
      return !prev;
    });
  }, [panelA]);

  return {
    datasets,
    presets,
    setPresets,
    error,
    setError,
    loading,
    activeDatasets,
    panelA,
    panelB,
    comparisonMode,
    addDataset,
    removeDataset,
    toggleDataset,
    updateDatasetPresetName,
    clearDatasetPreset,
    navigatePanel,
    goUp,
    toggleComparison,
    fetchOverview,
    fetchExampleDetail,
    fetchIterDetail,
    overviewCache,
    exampleDetailCache,
    iterDetailCache,
  };
}