frontend/src/rlm/store.ts

import { useState, useCallback, useEffect, useMemo } from "react";
import type {
  DatasetInfo,
  Preset,
  PanelNav,
  OverviewData,
  GepaIterData,
  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 [gepaIterCache, setGepaIterCache] = useState<Record<string, GepaIterData>>({});
  const [rlmDetailCache, setRlmDetailCache] = 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 fetchGepaIter = useCallback(
    async (dsId: string, gepaIter: number) => {
      const key = `${dsId}:${gepaIter}`;
      if (gepaIterCache[key]) return gepaIterCache[key];
      const data = (await api.getGepaIter(dsId, gepaIter)) as unknown as GepaIterData;
      setGepaIterCache((prev) => ({ ...prev, [key]: data }));
      return data;
    },
    [gepaIterCache]
  );

  const fetchRlmDetail = useCallback(
    async (dsId: string, gepaIter: number, rlmCallIdx: number, rlmIter: number) => {
      const key = `${dsId}:${gepaIter}:${rlmCallIdx}:${rlmIter}`;
      if (rlmDetailCache[key]) return rlmDetailCache[key];
      const data = (await api.getRlmIter(dsId, gepaIter, rlmCallIdx, rlmIter)) as unknown as RlmIterDetail;
      setRlmDetailCache((prev) => ({ ...prev, [key]: data }));
      return data;
    },
    [rlmDetailCache]
  );

  // 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_gepa_iters: result.n_gepa_iters,
          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 === 3) return { ...prev, level: 2, rlmIter: undefined };
      if (prev.level === 2) return { ...prev, level: 1, gepaIter: undefined, rlmCallIdx: 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,
    fetchGepaIter,
    fetchRlmDetail,
    overviewCache,
    gepaIterCache,
    rlmDetailCache,
  };
}