Pre_Seg_Server/src/store/useStore.ts

import { create } from 'zustand';

export interface Project {
  id: string;
  name: string;
  description?: string;
  status: 'pending' | 'parsing' | 'ready' | 'error';
  fps?: string;
  frames?: number;
  thumbnail?: string;
  thumbnail_url?: string;
  video_path?: string;
  source_type?: string;
  original_fps?: number;
  parse_fps?: number;
  createdAt?: string;
  updatedAt?: string;
}

export type AiModelId =
  | 'sam2.1_hiera_tiny'
  | 'sam2.1_hiera_small'
  | 'sam2.1_hiera_base_plus'
  | 'sam2.1_hiera_large';

export const DEFAULT_AI_MODEL_ID: AiModelId = 'sam2.1_hiera_tiny';
export const DEFAULT_BRUSH_SIZE = 24;
export const DEFAULT_ERASER_SIZE = 28;

export const SAM2_MODEL_OPTIONS: Array<{ id: AiModelId; label: string; shortLabel: string }> = [
  { id: 'sam2.1_hiera_tiny', label: 'SAM 2.1 Tiny', shortLabel: 'tiny' },
  { id: 'sam2.1_hiera_small', label: 'SAM 2.1 Small', shortLabel: 'small' },
  { id: 'sam2.1_hiera_base_plus', label: 'SAM 2.1 Base+', shortLabel: 'base+' },
  { id: 'sam2.1_hiera_large', label: 'SAM 2.1 Large', shortLabel: 'large' },
];

export interface Frame {
  id: string;
  projectId: string;
  index: number;
  url: string;
  width: number;
  height: number;
  timestamp?: string;
  timestampMs?: number;
  sourceFrameNumber?: number;
}

export interface Annotation {
  id: string;
  frameId: string;
  type: 'polygon' | 'rectangle' | 'circle' | 'point' | 'mask';
  points: number[];
  label: string;
  color: string;
  zIndex?: number;
  confidence?: number;
  metadata?: Record<string, unknown>;
}

export interface Mask {
  id: string;
  frameId: string;
  annotationId?: string;
  templateId?: string;
  classId?: string;
  className?: string;
  classZIndex?: number;
  classMaskId?: number;
  saveStatus?: 'draft' | 'saved' | 'dirty' | 'saving' | 'error';
  saved?: boolean;
  pathData: string;
  label: string;
  color: string;
  opacity?: number;
  segmentation?: number[][];
  points?: number[][];
  bbox?: [number, number, number, number];
  area?: number;
  metadata?: Record<string, unknown>;
}

export interface Template {
  id: string;
  name: string;
  description?: string;
  color?: string;
  z_index?: number;
  classes: TemplateClass[];
  rules?: TemplateRule[];
  createdAt?: string;
  updatedAt?: string;
}

export interface TemplateClass {
  id: string;
  name: string;
  color: string;
  zIndex: number;
  maskId?: number;
  category?: string;
  description?: string;
}

export interface TemplateRule {
  id: string;
  name: string;
  sourceKey: string;
  targetKey: string;
  operation: string;
}

export interface UserProfile {
  id: number;
  username: string;
  role: string;
  is_active?: number;
}

export interface AppState {
  // Auth
  isAuthenticated: boolean;
  token: string | null;
  currentUser: UserProfile | null;
  login: (token: string, user?: UserProfile | null) => void;
  setCurrentUser: (user: UserProfile | null) => void;
  logout: () => void;

  // Projects
  projects: Project[];
  currentProject: Project | null;
  setProjects: (projects: Project[]) => void;
  setCurrentProject: (project: Project | null) => void;
  addProject: (project: Project) => void;
  updateProject: (project: Project) => void;

  // Workspace
  activeModule: string;
  activeTool: string;
  aiModel: AiModelId;
  frames: Frame[];
  currentFrameIndex: number;
  annotations: Annotation[];
  masks: Mask[];
  selectedMaskIds: string[];
  maskPreviewOpacity: number;
  brushSize: number;
  eraserSize: number;
  maskHistory: Mask[][];
  maskFuture: Mask[][];
  setActiveModule: (module: string) => void;
  setActiveTool: (tool: string) => void;
  setAiModel: (model: AiModelId) => void;
  setFrames: (frames: Frame[]) => void;
  setCurrentFrame: (index: number) => void;
  addAnnotation: (annotation: Annotation) => void;
  addMask: (mask: Mask) => void;
  updateMask: (id: string, updates: Partial<Mask>) => void;
  setMasks: (masks: Mask[]) => void;
  setSelectedMaskIds: (ids: string[]) => void;
  setMaskPreviewOpacity: (opacity: number) => void;
  setBrushSize: (size: number) => void;
  setEraserSize: (size: number) => void;
  clearMasks: () => void;
  undoMasks: () => void;
  redoMasks: () => void;
  removeAnnotation: (id: string) => void;

  // Templates
  templates: Template[];
  activeTemplateId: string | null;
  activeClassId: string | null;
  activeClass: TemplateClass | null;
  setTemplates: (templates: Template[]) => void;
  setActiveTemplateId: (id: string | null) => void;
  setActiveClassId: (id: string | null) => void;
  setActiveClass: (templateClass: TemplateClass | null) => void;
  addTemplate: (template: Template) => void;
  updateTemplate: (template: Template) => void;
  removeTemplate: (id: string) => void;

  // UI
  isLoading: boolean;
  error: string | null;
  setLoading: (loading: boolean) => void;
  setError: (error: string | null) => void;
}

export const useStore = create<AppState>((set) => ({
  // Auth
  isAuthenticated: Boolean(localStorage.getItem('token')),
  token: localStorage.getItem('token'),
  currentUser: null,
  login: (token: string, user: UserProfile | null = null) => {
    localStorage.setItem('token', token);
    set({ isAuthenticated: true, token, currentUser: user });
  },
  setCurrentUser: (currentUser: UserProfile | null) => set({ currentUser }),
  logout: () => {
    localStorage.removeItem('token');
    set({
      isAuthenticated: false,
      token: null,
      currentUser: null,
      currentProject: null,
      projects: [],
      templates: [],
      frames: [],
      annotations: [],
      masks: [],
      selectedMaskIds: [],
      maskPreviewOpacity: 50,
      brushSize: DEFAULT_BRUSH_SIZE,
      eraserSize: DEFAULT_ERASER_SIZE,
      maskHistory: [],
      maskFuture: [],
      activeTemplateId: null,
      activeClassId: null,
      activeClass: null,
    });
  },

  // Projects
  projects: [],
  currentProject: null,
  setProjects: (projects: Project[]) => set({ projects }),
  setCurrentProject: (currentProject: Project | null) => set({ currentProject }),
  addProject: (project: Project) =>
    set((state) => ({ projects: [project, ...state.projects] })),
  updateProject: (project: Project) =>
    set((state) => ({
      projects: state.projects.map((p) => (p.id === project.id ? project : p)),
    })),

  // Workspace
  activeModule: 'dashboard',
  activeTool: 'move',
  aiModel: DEFAULT_AI_MODEL_ID,
  frames: [],
  currentFrameIndex: 0,
  annotations: [],
  masks: [],
  selectedMaskIds: [],
  maskPreviewOpacity: 50,
  brushSize: DEFAULT_BRUSH_SIZE,
  eraserSize: DEFAULT_ERASER_SIZE,
  maskHistory: [],
  maskFuture: [],
  setActiveModule: (activeModule: string) => set({ activeModule }),
  setActiveTool: (activeTool: string) => set({ activeTool }),
  setAiModel: (aiModel: AiModelId) => set({ aiModel }),
  setFrames: (frames: Frame[]) => set({ frames }),
  setCurrentFrame: (currentFrameIndex: number) => set({ currentFrameIndex }),
  addAnnotation: (annotation: Annotation) =>
    set((state) => ({ annotations: [...state.annotations, annotation] })),
  addMask: (mask: Mask) =>
    set((state) => ({
      masks: [...state.masks, mask],
      maskHistory: [...state.maskHistory, state.masks],
      maskFuture: [],
    })),
  updateMask: (id: string, updates: Partial<Mask>) =>
    set((state) => ({
      masks: state.masks.map((mask) => (mask.id === id ? { ...mask, ...updates } : mask)),
      maskHistory: [...state.maskHistory, state.masks],
      maskFuture: [],
    })),
  setMasks: (masks: Mask[]) =>
    set((state) => {
      const isInitialHydration = state.masks.length === 0
        && state.maskHistory.length === 0
        && state.maskFuture.length === 0;
      return {
        masks,
        maskHistory: isInitialHydration ? [] : [...state.maskHistory, state.masks],
        maskFuture: [],
      };
    }),
  setSelectedMaskIds: (selectedMaskIds: string[]) => set({ selectedMaskIds }),
  setMaskPreviewOpacity: (maskPreviewOpacity: number) => set({
    maskPreviewOpacity: Math.min(Math.max(maskPreviewOpacity, 10), 100),
  }),
  setBrushSize: (brushSize: number) => set({
    brushSize: Math.round(Math.min(Math.max(brushSize, 4), 96)),
  }),
  setEraserSize: (eraserSize: number) => set({
    eraserSize: Math.round(Math.min(Math.max(eraserSize, 4), 128)),
  }),
  clearMasks: () =>
    set((state) => ({
      masks: [],
      selectedMaskIds: [],
      maskHistory: [...state.maskHistory, state.masks],
      maskFuture: [],
    })),
  undoMasks: () =>
    set((state) => {
      if (state.maskHistory.length === 0) return state;
      const previous = state.maskHistory[state.maskHistory.length - 1];
      return {
        masks: previous,
        maskHistory: state.maskHistory.slice(0, -1),
        maskFuture: [state.masks, ...state.maskFuture],
      };
    }),
  redoMasks: () =>
    set((state) => {
      if (state.maskFuture.length === 0) return state;
      const [next, ...rest] = state.maskFuture;
      return {
        masks: next,
        maskHistory: [...state.maskHistory, state.masks],
        maskFuture: rest,
      };
    }),
  removeAnnotation: (id: string) =>
    set((state) => ({
      annotations: state.annotations.filter((a) => a.id !== id),
    })),

  // Templates
  templates: [],
  activeTemplateId: null,
  activeClassId: null,
  activeClass: null,
  setTemplates: (templates: Template[]) => set({ templates }),
  setActiveTemplateId: (activeTemplateId: string | null) => set({ activeTemplateId }),
  setActiveClassId: (activeClassId: string | null) => set({ activeClassId }),
  setActiveClass: (activeClass: TemplateClass | null) => set({
    activeClass,
    activeClassId: activeClass?.id || null,
  }),
  addTemplate: (template: Template) =>
    set((state) => ({ templates: [...state.templates, template] })),
  updateTemplate: (template: Template) =>
    set((state) => ({
      templates: state.templates.map((t) => (t.id === template.id ? template : t)),
    })),
  removeTemplate: (id: string) =>
    set((state) => ({
      templates: state.templates.filter((t) => t.id !== id),
    })),

  // UI
  isLoading: false,
  error: null,
  setLoading: (isLoading: boolean) => set({ isLoading }),
  setError: (error: string | null) => set({ error }),
}));