web-inspector/frontend/src/stores/useBatchInspectionStore.ts

import { create } from "zustand";
import type { AggregateScores } from "@/types/site-inspection";
import type {
  BatchPage,
  BatchInspectionPhase,
  BatchInspectionResult,
  SSEBatchPageComplete,
  SSEBatchAggregateUpdate,
} from "@/types/batch-inspection";

interface BatchInspectionState {
  batchInspectionId: string | null;
  name: string | null;
  status: BatchInspectionPhase;
  discoveredPages: BatchPage[];
  aggregateScores: AggregateScores | null;
  errorMessage: string | null;

  // Actions
  setBatchInspection: (id: string, name: string) => void;
  initFromApi: (data: BatchInspectionResult) => void;
  updatePageStatus: (
    pageUrl: string,
    status: BatchPage["status"],
    extra?: {
      inspection_id?: string;
      overall_score?: number;
      grade?: string;
    }
  ) => void;
  setPageComplete: (data: SSEBatchPageComplete) => void;
  updateAggregateScores: (data: SSEBatchAggregateUpdate) => void;
  setCompleted: (aggregateScores: AggregateScores) => void;
  setError: (message: string) => void;
  reset: () => void;
}

const initialState = {
  batchInspectionId: null,
  name: null,
  status: "idle" as BatchInspectionPhase,
  discoveredPages: [] as BatchPage[],
  aggregateScores: null,
  errorMessage: null,
};

export const useBatchInspectionStore = create<BatchInspectionState>(
  (set) => ({
    ...initialState,

    setBatchInspection: (id, name) =>
      set({
        ...initialState,
        batchInspectionId: id,
        name,
        status: "inspecting",
      }),

    initFromApi: (data) =>
      set((state) => {
        // Only init if store is idle (prevents overwriting live SSE data)
        if (state.status !== "idle") return state;
        return {
          batchInspectionId: data.batch_inspection_id,
          name: data.name,
          status: data.status as BatchInspectionPhase,
          discoveredPages: data.discovered_pages,
          aggregateScores: data.aggregate_scores,
        };
      }),

    updatePageStatus: (pageUrl, status, extra) =>
      set((state) => ({
        discoveredPages: state.discoveredPages.map((page) =>
          page.url === pageUrl
            ? {
                ...page,
                status,
                ...(extra?.inspection_id && {
                  inspection_id: extra.inspection_id,
                }),
                ...(extra?.overall_score !== undefined && {
                  overall_score: extra.overall_score,
                }),
                ...(extra?.grade && { grade: extra.grade }),
              }
            : page
        ),
      })),

    setPageComplete: (data) =>
      set((state) => ({
        discoveredPages: state.discoveredPages.map((page) =>
          page.url === data.page_url
            ? {
                ...page,
                status: "completed" as const,
                inspection_id: data.inspection_id,
                overall_score: data.overall_score,
                grade: data.grade,
              }
            : page
        ),
      })),

    updateAggregateScores: (data) =>
      set((state) => ({
        aggregateScores: state.aggregateScores
          ? {
              ...state.aggregateScores,
              pages_inspected: data.pages_inspected,
              pages_total: data.pages_total,
              overall_score: data.overall_score,
              grade: data.grade,
            }
          : {
              overall_score: data.overall_score,
              grade: data.grade,
              html_css: 0,
              accessibility: 0,
              seo: 0,
              performance_security: 0,
              total_issues: 0,
              pages_inspected: data.pages_inspected,
              pages_total: data.pages_total,
            },
      })),

    setCompleted: (aggregateScores) =>
      set({
        status: "completed",
        aggregateScores,
      }),

    setError: (message) =>
      set({
        status: "error",
        errorMessage: message,
      }),

    reset: () => set({ ...initialState }),
  })
);