emba-course-solver/app/src/state/appState.ts

import { useReducer, useMemo, useEffect, useRef, useCallback, useState } from 'react';
import { SPECIALIZATIONS } from '../data/specializations';
import { ELECTIVE_SETS } from '../data/electiveSets';
import type { OptimizationMode, AllocationResult } from '../data/types';
import { optimize } from '../solver/optimizer';
import type { SetAnalysis } from '../solver/decisionTree';
import type { WorkerRequest, WorkerResponse } from '../workers/decisionTree.worker';
import DecisionTreeWorker from '../workers/decisionTree.worker?worker';

const STORAGE_KEY = 'emba-solver-state';

export interface AppState {
  ranking: string[];
  mode: OptimizationMode;
  pinnedCourses: Record<string, string | null>; // setId -> courseId | null
}

type AppAction =
  | { type: 'reorder'; ranking: string[] }
  | { type: 'setMode'; mode: OptimizationMode }
  | { type: 'pinCourse'; setId: string; courseId: string }
  | { type: 'unpinCourse'; setId: string }
  | { type: 'clearAll' };

function reducer(state: AppState, action: AppAction): AppState {
  switch (action.type) {
    case 'reorder':
      return { ...state, ranking: action.ranking };
    case 'setMode':
      return { ...state, mode: action.mode };
    case 'pinCourse':
      return { ...state, pinnedCourses: { ...state.pinnedCourses, [action.setId]: action.courseId } };
    case 'unpinCourse': {
      const next = { ...state.pinnedCourses };
      delete next[action.setId];
      return { ...state, pinnedCourses: next };
    }
    case 'clearAll':
      return { ...state, pinnedCourses: {} };
  }
}

function defaultState(): AppState {
  return {
    ranking: SPECIALIZATIONS.map((s) => s.id),
    mode: 'maximize-count',
    pinnedCourses: {},
  };
}

function loadState(): AppState {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return defaultState();
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed.ranking) || parsed.ranking.length !== 14) return defaultState();
    if (!['maximize-count', 'priority-order'].includes(parsed.mode)) return defaultState();
    return {
      ranking: parsed.ranking,
      mode: parsed.mode,
      pinnedCourses: parsed.pinnedCourses ?? {},
    };
  } catch {
    return defaultState();
  }
}

export function useAppState() {
  const [state, dispatch] = useReducer(reducer, null, loadState);
  const [treeResults, setTreeResults] = useState<SetAnalysis[]>([]);
  const [treeLoading, setTreeLoading] = useState(false);
  const workerRef = useRef<Worker | null>(null);
  const debounceRef = useRef<ReturnType<typeof setTimeout>>();

  // Persist to localStorage
  useEffect(() => {
    localStorage.setItem(STORAGE_KEY, JSON.stringify(state));
  }, [state]);

  // Derive selected courses and open sets
  const selectedCourseIds = useMemo(
    () => Object.values(state.pinnedCourses).filter((v): v is string => v != null),
    [state.pinnedCourses],
  );

  const openSetIds = useMemo(
    () => ELECTIVE_SETS.map((s) => s.id).filter((id) => !state.pinnedCourses[id]),
    [state.pinnedCourses],
  );

  // Main-thread optimization (instant)
  const optimizationResult: AllocationResult = useMemo(
    () => optimize(selectedCourseIds, state.ranking, openSetIds, state.mode),
    [selectedCourseIds, state.ranking, openSetIds, state.mode],
  );

  // Web Worker decision tree (debounced)
  useEffect(() => {
    if (debounceRef.current) clearTimeout(debounceRef.current);

    if (openSetIds.length === 0) {
      setTreeResults([]);
      setTreeLoading(false);
      return;
    }

    setTreeLoading(true);

    debounceRef.current = setTimeout(() => {
      // Terminate previous worker if still running
      if (workerRef.current) workerRef.current.terminate();

      try {
        const worker = new DecisionTreeWorker();
        workerRef.current = worker;

        const progressResults: SetAnalysis[] = [];
        worker.onmessage = (e: MessageEvent<WorkerResponse>) => {
          if (e.data.type === 'setComplete' && e.data.analysis) {
            progressResults.push(e.data.analysis);
            setTreeResults([...progressResults]);
          } else if (e.data.type === 'allComplete' && e.data.analyses) {
            setTreeResults(e.data.analyses);
            setTreeLoading(false);
            worker.terminate();
            workerRef.current = null;
          }
        };

        const request: WorkerRequest = {
          pinnedCourseIds: selectedCourseIds,
          openSetIds,
          ranking: state.ranking,
          mode: state.mode,
        };
        worker.postMessage(request);
      } catch {
        // Web Worker not available (e.g., test env) — skip
        setTreeLoading(false);
      }
    }, 300);

    return () => {
      if (debounceRef.current) clearTimeout(debounceRef.current);
      if (workerRef.current) {
        workerRef.current.terminate();
        workerRef.current = null;
      }
    };
  }, [selectedCourseIds, openSetIds, state.ranking, state.mode]);

  const reorder = useCallback((ranking: string[]) => dispatch({ type: 'reorder', ranking }), []);
  const setMode = useCallback((mode: OptimizationMode) => dispatch({ type: 'setMode', mode }), []);
  const pinCourse = useCallback((setId: string, courseId: string) => dispatch({ type: 'pinCourse', setId, courseId }), []);
  const unpinCourse = useCallback((setId: string) => dispatch({ type: 'unpinCourse', setId }), []);
  const clearAll = useCallback(() => dispatch({ type: 'clearAll' }), []);

  return {
    state,
    optimizationResult,
    treeResults,
    treeLoading,
    openSetIds,
    selectedCourseIds,
    reorder,
    setMode,
    pinCourse,
    unpinCourse,
    clearAll,
  };
}