resolutionflow/frontend/src/store/aiChatStore.ts

import { create } from 'zustand'
import { AxiosError } from 'axios'
import { aiChatApi } from '@/api/aiChat'
import type {
  ChatMessage,
  InterviewPhase,
  TreeStructure,
} from '@/types'

interface TreeMetadata {
  name?: string
  description?: string
  tags?: string[]
  category_id?: string
}

interface AIChatState {
  // Session
  sessionId: string | null
  status: 'idle' | 'active' | 'completed' | 'abandoned'
  currentPhase: InterviewPhase
  flowType: 'troubleshooting' | 'procedural' | null

  // Conversation
  messages: ChatMessage[]
  isResponding: boolean

  // Progressive tree
  workingTree: TreeStructure | null
  treeMetadata: TreeMetadata | null

  // Final generation
  generatedTree: TreeStructure | null
  isGenerating: boolean
  importedTreeId: string | null

  // Error
  error: string | null

  // Actions
  startSession: (flowType: 'troubleshooting' | 'procedural') => Promise<void>
  sendMessage: (content: string) => Promise<void>
  generateTree: () => Promise<void>
  importToEditor: (params?: { name?: string; description?: string; category_id?: string; tags?: string[] }) => Promise<string>
  abandonSession: () => Promise<void>
  resumeSession: (sessionId: string) => Promise<void>
  reset: () => void
}

const initialState = {
  sessionId: null,
  status: 'idle' as const,
  currentPhase: 'scoping' as InterviewPhase,
  flowType: null,
  messages: [],
  isResponding: false,
  workingTree: null,
  treeMetadata: null,
  generatedTree: null,
  isGenerating: false,
  importedTreeId: null,
  error: null,
}

function extractErrorMessage(e: unknown, fallback: string): string {
  if (e instanceof AxiosError && e.response?.data?.detail) {
    const detail = e.response.data.detail
    return typeof detail === 'string' ? detail : detail.message || fallback
  }
  if (e instanceof Error) return e.message
  return fallback
}

export const useAIChatStore = create<AIChatState>((set, get) => ({
  ...initialState,

  startSession: async (flowType) => {
    set({ ...initialState, status: 'active', flowType, isResponding: true, error: null })

    try {
      const response = await aiChatApi.startSession(flowType)
      set({
        sessionId: response.session_id,
        currentPhase: response.current_phase,
        messages: [{
          role: 'assistant',
          content: response.greeting,
          timestamp: new Date().toISOString(),
        }],
        isResponding: false,
      })
    } catch (e: unknown) {
      set({ error: extractErrorMessage(e, 'Failed to start session'), isResponding: false, status: 'idle' })
    }
  },

  sendMessage: async (content) => {
    const { sessionId, messages, isResponding } = get()
    if (!sessionId || isResponding) return

    const userMessage: ChatMessage = {
      role: 'user',
      content,
      timestamp: new Date().toISOString(),
    }

    set({
      messages: [...messages, userMessage],
      isResponding: true,
      error: null,
    })

    try {
      const response = await aiChatApi.sendMessage(sessionId, content)
      const aiMessage: ChatMessage = {
        role: 'assistant',
        content: response.content,
        timestamp: new Date().toISOString(),
      }

      set((state) => ({
        messages: [...state.messages, aiMessage],
        currentPhase: response.current_phase,
        workingTree: (response.working_tree as TreeStructure | null) ?? state.workingTree,
        treeMetadata: (response.tree_metadata as TreeMetadata | null) ?? state.treeMetadata,
        isResponding: false,
      }))
    } catch (e: unknown) {
      set({ error: extractErrorMessage(e, 'Failed to send message'), isResponding: false })
    }
  },

  generateTree: async () => {
    const { sessionId, isGenerating } = get()
    if (!sessionId || isGenerating) return

    set({ isGenerating: true, error: null })

    try {
      const response = await aiChatApi.generateTree(sessionId)
      set({
        generatedTree: response.tree_structure as unknown as TreeStructure,
        workingTree: response.tree_structure as unknown as TreeStructure,
        treeMetadata: response.tree_metadata as TreeMetadata,
        status: 'completed',
        isGenerating: false,
      })
    } catch (e: unknown) {
      set({ error: extractErrorMessage(e, 'Failed to generate tree'), isGenerating: false })
    }
  },

  importToEditor: async (params) => {
    const { sessionId } = get()
    if (!sessionId) throw new Error('No active session')

    const response = await aiChatApi.importTree(sessionId, params)
    set({ importedTreeId: response.tree_id })
    return response.tree_id
  },

  abandonSession: async () => {
    const { sessionId } = get()
    if (!sessionId) return

    try {
      await aiChatApi.abandonSession(sessionId)
    } catch {
      // Best effort — session may have already expired
    }
    set({ ...initialState })
  },

  resumeSession: async (sessionId) => {
    set({ isResponding: true, error: null })

    try {
      const session = await aiChatApi.getSession(sessionId)
      set({
        sessionId: session.session_id,
        status: session.status === 'active' ? 'active' : (session.status as 'completed' | 'abandoned'),
        currentPhase: session.current_phase,
        flowType: session.flow_type,
        messages: session.conversation_history as ChatMessage[],
        workingTree: session.working_tree as TreeStructure | null,
        treeMetadata: session.tree_metadata as TreeMetadata | null,
        generatedTree: session.generated_tree as TreeStructure | null,
        isResponding: false,
      })
    } catch (e: unknown) {
      set({ error: extractErrorMessage(e, 'Failed to resume session'), isResponding: false })
    }
  },

  reset: () => set({ ...initialState }),
}))