import { create } from 'zustand'

import type { Source } from '../../lib/types'

type Phase = 'idle' | 'researching' | 'answering' | 'error'

type ChatStreamState = {
  phase: Phase
  activeQuery: string
  text: string
  sources: Source[]
  selectedSource: Source | null
  error: string | null
  start: (query: string) => void
  setPhase: (phase: Phase) => void
  appendText: (content: string) => void
  setSources: (sources: Source[]) => void
  selectSource: (source: Source | null) => void
  fail: (message: string) => void
  reset: () => void
}

const initial = {
  phase: 'idle' as Phase,
  activeQuery: '',
  text: '',
  sources: [] as Source[],
  selectedSource: null as Source | null,
  error: null as string | null,
}

export const useChatStreamStore = create<ChatStreamState>((set) => ({
  ...initial,
  start: (query) => set({ ...initial, phase: 'researching', activeQuery: query }),
  setPhase: (phase) => set({ phase }),
  appendText: (content) => set((state) => ({ text: state.text + content })),
  setSources: (sources) => set({ sources, selectedSource: sources[0] ?? null }),
  selectSource: (selectedSource) => set({ selectedSource }),
  fail: (error) => set({ phase: 'error', error }),
  reset: () => set(initial),
}))
