resolutionflow/frontend/src/store/aiFlowBuilderStore.ts

import { create } from 'zustand'
import { aiBuilderApi } from '@/api/aiBuilder'
import type { AIQuotaStatus, AIBranch, AIAssembleResponse, AIWizardPhase } from '@/types'

interface AIFlowBuilderState {
  // Wizard state
  phase: AIWizardPhase
  conversationId: string | null
  metadata: {
    flow_type: 'troubleshooting' | 'procedural'
    name: string
    description: string
    environment_tags: string[]
    category_id: string | null
  }

  // Stage 2
  suggestedBranches: AIBranch[]
  selectedBranches: AIBranch[]

  // Stage 3
  currentBranchIndex: number

  // Stage 4
  assembledTree: AIAssembleResponse | null

  // Quota
  quota: AIQuotaStatus | null

  // UI state
  isLoading: boolean
  error: string | null

  // Actions
  loadQuota: () => Promise<void>
  setMetadata: (metadata: Partial<AIFlowBuilderState['metadata']>) => void
  start: () => Promise<void>
  scaffold: () => Promise<void>
  selectBranches: (branches: AIBranch[]) => void
  generateBranchDetail: (branchName: string) => Promise<void>
  assemble: () => Promise<void>
  reset: () => void
  setPhase: (phase: AIWizardPhase) => void
  setError: (error: string | null) => void
}

const initialMetadata = {
  flow_type: 'troubleshooting' as const,
  name: '',
  description: '',
  environment_tags: [] as string[],
  category_id: null as string | null,
}

export const useAIFlowBuilderStore = create<AIFlowBuilderState>()((set, get) => ({
  phase: 'foundation',
  conversationId: null,
  metadata: { ...initialMetadata },
  suggestedBranches: [],
  selectedBranches: [],
  currentBranchIndex: 0,
  assembledTree: null,
  quota: null,
  isLoading: false,
  error: null,

  loadQuota: async () => {
    try {
      const quota = await aiBuilderApi.getQuota()
      set({ quota })
    } catch {
      // Silently fail — quota display is optional
    }
  },

  setMetadata: (metadata) => {
    set((state) => ({
      metadata: { ...state.metadata, ...metadata },
    }))
  },

  start: async () => {
    const { metadata } = get()
    set({ isLoading: true, error: null })
    try {
      const response = await aiBuilderApi.start({
        flow_type: metadata.flow_type,
        name: metadata.name,
        description: metadata.description,
        environment_tags: metadata.environment_tags,
        category_id: metadata.category_id ?? undefined,
      })
      set({
        conversationId: response.conversation_id,
        phase: 'scaffolding',
        isLoading: false,
      })
    } catch (err) {
      const message = _extractError(err)
      set({ error: message, isLoading: false })
    }
  },

  scaffold: async () => {
    const { conversationId } = get()
    if (!conversationId) return
    set({ isLoading: true, error: null, phase: 'generating' })
    try {
      const response = await aiBuilderApi.scaffold(conversationId)
      const branches: AIBranch[] = response.branches.map((b) => ({
        name: b.name,
        description: b.description,
      }))
      set({
        suggestedBranches: branches,
        selectedBranches: branches,
        phase: 'scaffolding',
        isLoading: false,
      })
    } catch (err) {
      const message = _extractError(err)
      set({ error: message, phase: 'error', isLoading: false })
    }
  },

  selectBranches: (branches) => {
    set({ selectedBranches: branches })
  },

  generateBranchDetail: async (branchName) => {
    const { conversationId, selectedBranches } = get()
    if (!conversationId) return
    set({ isLoading: true, error: null, phase: 'generating' })
    try {
      const response = await aiBuilderApi.branchDetail(conversationId, branchName)
      const updatedBranches = selectedBranches.map((b) =>
        b.name === branchName ? { ...b, steps: response.steps } : b
      )
      // Advance to the next branch that still needs detail
      const nextIndex = updatedBranches.findIndex((b) => !b.steps)
      const currentBranchIndex = nextIndex !== -1 ? nextIndex : updatedBranches.findIndex((b) => b.name === branchName)
      set({
        selectedBranches: updatedBranches,
        currentBranchIndex,
        phase: 'detailing',
        isLoading: false,
      })
    } catch (err) {
      const message = _extractError(err)
      set({ error: message, phase: 'error', isLoading: false })
    }
  },

  assemble: async () => {
    const { conversationId, selectedBranches } = get()
    if (!conversationId) return
    set({ isLoading: true, error: null })
    try {
      const response = await aiBuilderApi.assemble(
        conversationId,
        selectedBranches.map((b) => ({
          name: b.name,
          description: b.description,
          steps: b.steps,
        }))
      )
      set({
        assembledTree: response,
        phase: 'reviewing',
        isLoading: false,
      })
    } catch (err) {
      const message = _extractError(err)
      set({ error: message, phase: 'error', isLoading: false })
    }
  },

  reset: () => {
    set({
      phase: 'foundation',
      conversationId: null,
      metadata: { ...initialMetadata },
      suggestedBranches: [],
      selectedBranches: [],
      currentBranchIndex: 0,
      assembledTree: null,
      isLoading: false,
      error: null,
    })
  },

  setPhase: (phase) => set({ phase }),
  setError: (error) => set({ error }),
}))

function _extractError(err: unknown): string {
  if (err && typeof err === 'object' && 'response' in err) {
    const axiosErr = err as { response?: { data?: { detail?: string | { message?: string } } } }
    const detail = axiosErr.response?.data?.detail
    if (typeof detail === 'string') return detail
    if (detail && typeof detail === 'object' && 'message' in detail) return detail.message ?? 'Unknown error'
  }
  if (err instanceof Error) return err.message
  return 'An unexpected error occurred'
}