frontend/src/adaevolve/store.ts

import { useCallback, useState } from "react";
import type {
  DatasetInfo,
  IterationDetail,
  Preset,
  ViewMode,
} from "./types";
import * as api from "./api";

export interface AppState {
  datasets: DatasetInfo[];
  presets: Preset[];
  selectedDatasetId: string | null;
  selectedIterationIdx: number | null;
  iterationDetail: IterationDetail | null;
  viewMode: ViewMode;
  loading: boolean;
  error: string | null;
  filterAdaptation: string; // "" means all
}

const initialState: AppState = {
  datasets: [],
  presets: [],
  selectedDatasetId: null,
  selectedIterationIdx: null,
  iterationDetail: null,
  viewMode: "timeline",
  loading: false,
  error: null,
  filterAdaptation: "",
};

export function useAppState() {
  const [state, setState] = useState<AppState>(initialState);

  const setError = useCallback((error: string | null) => {
    setState((s) => ({ ...s, error }));
  }, []);

  const loadDataset = useCallback(async (repo: string, split = "train") => {
    setState((s) => ({ ...s, loading: true, error: null }));
    try {
      const ds = await api.loadDataset(repo, split);
      setState((s) => {
        const exists = s.datasets.find((d) => d.id === ds.id);
        const datasets = exists
          ? s.datasets.map((d) => (d.id === ds.id ? ds : d))
          : [...s.datasets, ds];
        const selectedDatasetId = ds.id;
        return {
          ...s,
          datasets,
          selectedDatasetId,
          loading: false,
        };
      });
    } catch (e: unknown) {
      setState((s) => ({
        ...s,
        loading: false,
        error: e instanceof Error ? e.message : String(e),
      }));
    }
  }, []);

  const unloadDataset = useCallback(async (dsId: string) => {
    try {
      await api.unloadDataset(dsId);
      setState((s) => {
        const datasets = s.datasets.filter((d) => d.id !== dsId);
        return {
          ...s,
          datasets,
          selectedDatasetId: s.selectedDatasetId === dsId ? null : s.selectedDatasetId,
          selectedIterationIdx: s.selectedDatasetId === dsId ? null : s.selectedIterationIdx,
          iterationDetail: s.selectedDatasetId === dsId ? null : s.iterationDetail,
        };
      });
    } catch (e: unknown) {
      setState((s) => ({
        ...s,
        error: e instanceof Error ? e.message : String(e),
      }));
    }
  }, []);

  const selectDataset = useCallback((dsId: string | null) => {
    setState((s) => ({
      ...s,
      selectedDatasetId: dsId,
      selectedIterationIdx: null,
      iterationDetail: null,
      viewMode: "timeline",
    }));
  }, []);

  const selectIteration = useCallback(async (dsId: string, idx: number) => {
    setState((s) => ({ ...s, loading: true, error: null }));
    try {
      const detail = await api.getIterationDetail(dsId, idx);
      setState((s) => ({
        ...s,
        selectedDatasetId: dsId,
        selectedIterationIdx: idx,
        iterationDetail: detail,
        viewMode: "detail",
        loading: false,
      }));
    } catch (e: unknown) {
      setState((s) => ({
        ...s,
        loading: false,
        error: e instanceof Error ? e.message : String(e),
      }));
    }
  }, []);

  const backToTimeline = useCallback(() => {
    setState((s) => ({
      ...s,
      selectedIterationIdx: null,
      iterationDetail: null,
      viewMode: "timeline",
    }));
  }, []);

  const loadPresets = useCallback(async () => {
    try {
      const presets = await api.listPresets();
      setState((s) => ({ ...s, presets }));
    } catch {
      // Presets might not exist yet
    }
  }, []);

  const createPreset = useCallback(
    async (name: string, repo: string, split = "train") => {
      const preset = await api.createPreset(name, repo, split);
      setState((s) => ({ ...s, presets: [...s.presets, preset] }));
    },
    []
  );

  const deletePreset = useCallback(async (id: string) => {
    await api.deletePreset(id);
    setState((s) => ({ ...s, presets: s.presets.filter((p) => p.id !== id) }));
  }, []);

  const loadPreset = useCallback(
    async (preset: Preset) => {
      await loadDataset(preset.repo, preset.split);
    },
    [loadDataset]
  );

  const setViewMode = useCallback((viewMode: ViewMode) => {
    setState((s) => ({ ...s, viewMode }));
  }, []);

  const setFilterAdaptation = useCallback((filterAdaptation: string) => {
    setState((s) => ({ ...s, filterAdaptation }));
  }, []);

  return {
    state,
    loadDataset,
    unloadDataset,
    selectDataset,
    selectIteration,
    backToTimeline,
    loadPresets,
    createPreset,
    deletePreset,
    loadPreset,
    setViewMode,
    setFilterAdaptation,
    setError,
  };
}