frontend/src/store/gridStore.ts

import { create } from 'zustand';
import type {
  GridState,
  SearchResult,
  PlanResult,
  ComparisonResult,
  Algorithm,
  SearchStep,
  GridConfig,
  Position,
} from '../types';
import { generateGrid, createPlan, compareAlgorithms, findPath } from '../api/client';

// Path with color for visualization
export interface ColoredPath {
  storeId: number;
  destinationId: number;
  path: Position[];
  color: string;
}

interface GridStore {
  // Grid state
  grid: GridState | null;
  initialState: string;
  traffic: string;

  // Search state
  searchResult: SearchResult | null;
  planResult: PlanResult | null;
  comparisonResults: ComparisonResult[] | null;

  // Visualization state
  currentStep: number;
  steps: SearchStep[];
  isPlaying: boolean;
  playbackSpeed: number; // ms per step

  // Plan visualization
  showPlanPaths: boolean;
  planPaths: ColoredPath[];

  // UI state
  selectedAlgorithm: Algorithm;
  isLoading: boolean;
  error: string | null;

  // Actions
  setGrid: (grid: GridState, initialState: string, traffic: string) => void;
  generateNewGrid: (config?: GridConfig) => Promise<void>;
  runSearch: (start: { x: number; y: number }, goal: { x: number; y: number }) => Promise<void>;
  runPlan: (visualize?: boolean) => Promise<void>;
  runComparison: () => Promise<void>;
  setAlgorithm: (algorithm: Algorithm) => void;
  setCurrentStep: (step: number) => void;
  play: () => void;
  pause: () => void;
  setSpeed: (speed: number) => void;
  reset: () => void;
  nextStep: () => void;
  prevStep: () => void;
  setShowPlanPaths: (show: boolean) => void;
  clearPlanPaths: () => void;
}

// Colors for different delivery paths
const PATH_COLORS = [
  '#f97316', // orange
  '#06b6d4', // cyan
  '#ec4899', // pink
  '#84cc16', // lime
  '#a855f7', // purple
  '#14b8a6', // teal
  '#f43f5e', // rose
  '#eab308', // yellow
  '#6366f1', // indigo
  '#22c55e', // green
];

export const useGridStore = create<GridStore>((set, get) => ({
  // Initial state
  grid: null,
  initialState: '',
  traffic: '',
  searchResult: null,
  planResult: null,
  comparisonResults: null,
  currentStep: 0,
  steps: [],
  isPlaying: false,
  playbackSpeed: 100,
  showPlanPaths: false,
  planPaths: [],
  selectedAlgorithm: 'BF',
  isLoading: false,
  error: null,

  // Actions
  setGrid: (grid, initialState, traffic) => {
    set({
      grid,
      initialState,
      traffic,
      searchResult: null,
      planResult: null,
      comparisonResults: null,
      currentStep: 0,
      steps: [],
      isPlaying: false,
      showPlanPaths: false,
      planPaths: [],
      error: null,
    });
  },

  generateNewGrid: async (config = {}) => {
    set({ isLoading: true, error: null });
    try {
      const result = await generateGrid(config);
      set({
        grid: result.parsed,
        initialState: result.initialState,
        traffic: result.traffic,
        searchResult: null,
        planResult: null,
        comparisonResults: null,
        currentStep: 0,
        steps: [],
        isPlaying: false,
        showPlanPaths: false,
        planPaths: [],
        isLoading: false,
      });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Failed to generate grid',
        isLoading: false,
      });
    }
  },

  runSearch: async (start, goal) => {
    const { grid, selectedAlgorithm } = get();
    if (!grid) return;

    set({ isLoading: true, error: null, showPlanPaths: false, planPaths: [] });
    try {
      const result = await findPath(
        grid.width,
        grid.height,
        start,
        goal,
        grid.segments,
        grid.tunnels,
        selectedAlgorithm
      );
      set({
        searchResult: result,
        steps: result.steps || [],
        currentStep: 0,
        isPlaying: false,
        isLoading: false,
      });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Search failed',
        isLoading: false,
      });
    }
  },

  runPlan: async (visualize = false) => {
    const { initialState, traffic, selectedAlgorithm } = get();
    if (!initialState) return;

    set({ isLoading: true, error: null });
    try {
      const result = await createPlan(initialState, traffic, selectedAlgorithm, visualize);

      // Extract paths for visualization with different colors
      const coloredPaths: ColoredPath[] = result.assignments.map((assignment, index) => ({
        storeId: assignment.store_id,
        destinationId: assignment.destination_id,
        path: assignment.path.path,
        color: PATH_COLORS[index % PATH_COLORS.length],
      }));

      set({
        planResult: result,
        planPaths: coloredPaths,
        showPlanPaths: true,
        isLoading: false,
      });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Planning failed',
        isLoading: false,
      });
    }
  },

  runComparison: async () => {
    const { initialState, traffic } = get();
    if (!initialState) return;

    set({ isLoading: true, error: null });
    try {
      const result = await compareAlgorithms(initialState, traffic);
      set({
        comparisonResults: result.comparisons,
        isLoading: false,
      });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Comparison failed',
        isLoading: false,
      });
    }
  },

  setAlgorithm: (algorithm) => {
    set({ selectedAlgorithm: algorithm });
  },

  setCurrentStep: (step) => {
    const { steps } = get();
    set({ currentStep: Math.max(0, Math.min(step, steps.length - 1)) });
  },

  play: () => {
    set({ isPlaying: true });
  },

  pause: () => {
    set({ isPlaying: false });
  },

  setSpeed: (speed) => {
    set({ playbackSpeed: speed });
  },

  reset: () => {
    set({
      currentStep: 0,
      isPlaying: false,
    });
  },

  nextStep: () => {
    const { currentStep, steps } = get();
    if (currentStep < steps.length - 1) {
      set({ currentStep: currentStep + 1 });
    } else {
      set({ isPlaying: false });
    }
  },

  prevStep: () => {
    const { currentStep } = get();
    if (currentStep > 0) {
      set({ currentStep: currentStep - 1 });
    }
  },

  setShowPlanPaths: (show) => {
    set({ showPlanPaths: show });
  },

  clearPlanPaths: () => {
    set({ showPlanPaths: false, planPaths: [] });
  },
}));