import { create } from "zustand";
import type { DetectResponse, VisualizeResponse, ModelInfo } from "@/types/inference";

export type AudioSource =
  | { kind: "file"; file: File; objectUrl: string }
  | { kind: "sample"; sampleId: string; url: string }
  | { kind: "recording"; blob: Blob; objectUrl: string }
  | null;

interface InferenceState {
  audio: AudioSource;
  setAudio: (a: AudioSource) => void;

  selectedModels: string[];
  setSelectedModels: (m: string[]) => void;
  toggleModel: (id: string) => void;

  models: ModelInfo[];
  setModels: (m: ModelInfo[]) => void;

  isInferring: boolean;
  setIsInferring: (v: boolean) => void;

  result: DetectResponse | null;
  setResult: (r: DetectResponse | null) => void;

  viz: VisualizeResponse | null;
  setViz: (v: VisualizeResponse | null) => void;

  error: string | null;
  setError: (e: string | null) => void;

  reset: () => void;
}

export const useInferenceStore = create<InferenceState>((set) => ({
  audio: null,
  setAudio: (a) => set({ audio: a, result: null, viz: null, error: null }),

  selectedModels: [],
  setSelectedModels: (m) => set({ selectedModels: m }),
  toggleModel: (id) =>
    set((s) => ({
      selectedModels: s.selectedModels.includes(id)
        ? s.selectedModels.filter((x) => x !== id)
        : [...s.selectedModels, id],
    })),

  models: [],
  setModels: (m) => set({ models: m }),

  isInferring: false,
  setIsInferring: (v) => set({ isInferring: v }),

  result: null,
  setResult: (r) => set({ result: r }),

  viz: null,
  setViz: (v) => set({ viz: v }),

  error: null,
  setError: (e) => set({ error: e }),

  reset: () =>
    set({
      audio: null,
      result: null,
      viz: null,
      error: null,
      isInferring: false,
    }),
}));