import { create } from "zustand";
import { persist } from "zustand/middleware";
import type { GenerateResponse, MatrixGenerateResponse, GenerationProgress } from "../types";

interface GenerationState {
  currentGeneration: GenerateResponse | MatrixGenerateResponse | null;
  progress: GenerationProgress;
  isGenerating: boolean;
  error: string | null;
  generationStartTime: number | null;
  
  setCurrentGeneration: (ad: GenerateResponse | MatrixGenerateResponse | null) => void;
  setProgress: (progress: GenerationProgress) => void;
  setIsGenerating: (isGenerating: boolean) => void;
  setError: (error: string | null) => void;
  setGenerationStartTime: (time: number | null) => void;
  reset: () => void;
}

const initialState = {
  currentGeneration: null as GenerateResponse | MatrixGenerateResponse | null,
  progress: {
    step: "idle" as const,
    progress: 0,
  },
  isGenerating: false,
  error: null as string | null,
  generationStartTime: null as number | null,
};

// Storage key for session persistence
const STORAGE_KEY = "generation-state";

export const useGenerationStore = create<GenerationState>()(
  persist(
    (set) => ({
      ...initialState,
      
      setCurrentGeneration: (ad) => set({ currentGeneration: ad }),
      
      setProgress: (progress) => set({ progress }),
      
      setIsGenerating: (isGenerating) => set({ isGenerating }),
      
      setError: (error) => set({ error }),
      
      setGenerationStartTime: (time) => set({ generationStartTime: time }),
      
      reset: () => set(initialState),
    }),
    {
      name: STORAGE_KEY,
      // Only persist when generating to avoid stale data
      partialize: (state) => ({
        progress: state.progress,
        isGenerating: state.isGenerating,
        generationStartTime: state.generationStartTime,
        currentGeneration: state.isGenerating ? state.currentGeneration : null,
      }),
    }
  )
);