// Chat context for LLM assistant state management

import { createContext, useContext, useState, useCallback, useEffect, ReactNode, useRef } from 'react'
import * as chatApi from '@/api/chat'
import type { ChatSession, ChatMessage, PageContext, LlmModel } from '@/types/chat'
import { useAuth } from './AuthContext'

interface ChatContextValue {
  // UI State
  isOpen: boolean
  isMinimized: boolean
  toggleChat: () => void
  minimizeChat: () => void
  expandChat: () => void
  closeChat: () => void

  // Session Management
  sessions: ChatSession[]
  currentSession: ChatSession | null
  isLoadingSessions: boolean
  createSession: (workspaceId?: string) => Promise<ChatSession>
  switchSession: (sessionId: string) => Promise<void>
  deleteSession: (sessionId: string) => Promise<void>
  refreshSessions: () => Promise<void>

  // Messages
  messages: ChatMessage[]
  isStreaming: boolean
  streamingContent: string
  sendMessage: (content: string) => Promise<void>
  clearMessages: () => Promise<void>

  // Models
  models: LlmModel[]
  selectedModelId: string | null
  setSelectedModelId: (modelId: string) => void
  loadModels: () => Promise<void>

  // Page Context
  currentPageContext: PageContext | null
  setCurrentPageContext: (context: PageContext | null) => void
}

const ChatContext = createContext<ChatContextValue | null>(null)

interface ChatProviderProps {
  children: ReactNode
}

export function ChatProvider({ children }: ChatProviderProps) {
  const { isAuthenticated } = useAuth()

  // UI State
  const [isOpen, setIsOpen] = useState(false)
  const [isMinimized, setIsMinimized] = useState(true)

  // Session State
  const [sessions, setSessions] = useState<ChatSession[]>([])
  const [currentSession, setCurrentSession] = useState<ChatSession | null>(null)
  const [isLoadingSessions, setIsLoadingSessions] = useState(false)

  // Message State
  const [messages, setMessages] = useState<ChatMessage[]>([])
  const [isStreaming, setIsStreaming] = useState(false)
  const [streamingContent, setStreamingContent] = useState('')

  // Model State
  const [models, setModels] = useState<LlmModel[]>([])
  const [selectedModelId, setSelectedModelId] = useState<string | null>(null)

  // Page Context
  const [currentPageContext, setCurrentPageContext] = useState<PageContext | null>(null)

  // Abort controller for streaming
  const abortControllerRef = useRef<AbortController | null>(null)

  // UI Actions
  const toggleChat = useCallback(() => {
    if (isMinimized) {
      setIsMinimized(false)
      setIsOpen(true)
    } else {
      setIsMinimized(true)
    }
  }, [isMinimized])

  const minimizeChat = useCallback(() => {
    setIsMinimized(true)
  }, [])

  const expandChat = useCallback(() => {
    setIsMinimized(false)
    setIsOpen(true)
  }, [])

  const closeChat = useCallback(() => {
    setIsOpen(false)
    setIsMinimized(true)
  }, [])

  // Session Actions
  const refreshSessions = useCallback(async () => {
    if (!isAuthenticated) return

    setIsLoadingSessions(true)
    try {
      const response = await chatApi.listSessions(undefined, 50)
      setSessions(response.sessions || [])
    } catch (error) {
      console.error('Failed to load sessions:', error)
    } finally {
      setIsLoadingSessions(false)
    }
  }, [isAuthenticated])

  const createSession = useCallback(async (workspaceId?: string): Promise<ChatSession> => {
    const session = await chatApi.createSession({
      workspace_id: workspaceId,
      model_id: selectedModelId || undefined,
    })
    setSessions(prev => [session, ...prev])
    setCurrentSession(session)
    setMessages([])
    return session
  }, [selectedModelId])

  const switchSession = useCallback(async (sessionId: string) => {
    try {
      const session = await chatApi.getSession(sessionId)
      setCurrentSession(session)
      setMessages(session.messages || [])
    } catch (error) {
      console.error('Failed to switch session:', error)
      throw error
    }
  }, [])

  const deleteSession = useCallback(async (sessionId: string) => {
    await chatApi.deleteSession(sessionId)
    setSessions(prev => prev.filter(s => s.id !== sessionId))
    if (currentSession?.id === sessionId) {
      setCurrentSession(null)
      setMessages([])
    }
  }, [currentSession?.id])

  // Message Actions
  const sendMessage = useCallback(async (content: string) => {
    if (!currentSession) {
      // Create a new session first
      const session = await createSession()
      setCurrentSession(session)
    }

    const sessionId = currentSession?.id
    if (!sessionId) return

    // Add user message optimistically
    const userMessage: ChatMessage = {
      id: `temp-${Date.now()}`,
      role: 'user',
      content,
      created_at: Date.now() / 1000,
      page_context: currentPageContext || undefined,
    }
    setMessages(prev => [...prev, userMessage])

    // Start streaming
    setIsStreaming(true)
    setStreamingContent('')

    try {
      await chatApi.streamMessage(
        sessionId,
        {
          content,
          page_context: currentPageContext || undefined,
        },
        (chunk) => {
          if (chunk.delta) {
            setStreamingContent(prev => prev + chunk.delta)
          }
          if (chunk.finish_reason && chunk.message_id) {
            // Stream complete, add the full message
            const assistantMessage: ChatMessage = {
              id: chunk.message_id,
              role: 'assistant',
              content: '', // Will be set below
              created_at: Date.now() / 1000,
            }
            setMessages(prev => {
              // Get current streaming content
              const fullContent = prev.length > 0 ? streamingContent : ''
              return [...prev, { ...assistantMessage, content: fullContent }]
            })
          }
        },
        (error) => {
          console.error('Stream error:', error)
        }
      )

      // After streaming completes, refresh the session to get proper message IDs
      const updatedSession = await chatApi.getSession(sessionId)
      setMessages(updatedSession.messages || [])
      setCurrentSession(updatedSession)

      // Update session in list
      setSessions(prev => prev.map(s =>
        s.id === sessionId ? { ...s, title: updatedSession.title, updated_at: updatedSession.updated_at } : s
      ))
    } catch (error) {
      console.error('Failed to send message:', error)
      // Remove optimistic user message on error
      setMessages(prev => prev.filter(m => m.id !== userMessage.id))
    } finally {
      setIsStreaming(false)
      setStreamingContent('')
    }
  }, [currentSession, currentPageContext, createSession, streamingContent])

  const clearMessages = useCallback(async () => {
    if (!currentSession) return

    await chatApi.clearSessionMessages(currentSession.id)
    setMessages([])
  }, [currentSession])

  // Model Actions
  const loadModels = useCallback(async () => {
    try {
      const response = await chatApi.listModels()
      setModels(response.models || [])
      // Select first model if none selected
      if (!selectedModelId && response.models?.length > 0) {
        setSelectedModelId(response.models[0].id)
      }
    } catch (error) {
      console.error('Failed to load models:', error)
    }
  }, [selectedModelId])

  // Load sessions and models on mount when authenticated
  useEffect(() => {
    if (isAuthenticated) {
      refreshSessions()
      loadModels()
    }
  }, [isAuthenticated, refreshSessions, loadModels])

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (abortControllerRef.current) {
        abortControllerRef.current.abort()
      }
    }
  }, [])

  const value: ChatContextValue = {
    // UI State
    isOpen,
    isMinimized,
    toggleChat,
    minimizeChat,
    expandChat,
    closeChat,

    // Session Management
    sessions,
    currentSession,
    isLoadingSessions,
    createSession,
    switchSession,
    deleteSession,
    refreshSessions,

    // Messages
    messages,
    isStreaming,
    streamingContent,
    sendMessage,
    clearMessages,

    // Models
    models,
    selectedModelId,
    setSelectedModelId,
    loadModels,

    // Page Context
    currentPageContext,
    setCurrentPageContext,
  }

  return (
    <ChatContext.Provider value={value}>
      {children}
    </ChatContext.Provider>
  )
}

export function useChat(): ChatContextValue {
  const context = useContext(ChatContext)
  if (!context) {
    throw new Error('useChat must be used within a ChatProvider')
  }
  return context
}