import { useState } from 'react'

const API_BASE = '/api'

export interface Citation {
  id: string
  text: string
  source_file: string
  source_directory: string
  section: string | null
  date: string | null
  tags: string[]
  citation: string
}

export interface StreamCallbacks {
  onChunk: (chunk: string) => void
  onCitations?: (citations: Citation[]) => void
  onError?: (error: string) => void
  onDone?: () => void
}

export function useChatStream() {
  const [sessionId, setSessionId] = useState<string | null>(null)

  const sendMessage = async (
    message: string,
    callbacks: StreamCallbacks
  ): Promise<void> => {
    const { onChunk, onCitations, onError, onDone } = callbacks

    const response = await fetch(`${API_BASE}/chat`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        message,
        session_id: sessionId,
        stream: true,
        use_rag: true,
      }),
    })

    if (!response.ok) {
      throw new Error(`HTTP error! status: ${response.status}`)
    }

    // Get session ID from response headers if present
    const newSessionId = response.headers.get('X-Session-ID')
    if (newSessionId) {
      setSessionId(newSessionId)
    }

    const reader = response.body?.getReader()
    const decoder = new TextDecoder()

    if (!reader) {
      throw new Error('No response body')
    }

    try {
      while (true) {
        const { done, value } = await reader.read()
        if (done) break

        const chunk = decoder.decode(value, { stream: true })
        const lines = chunk.split('\n')

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6)
            if (data === '[DONE]') {
              onDone?.()
              return
            }
            try {
              const parsed = JSON.parse(data)
              switch (parsed.type) {
                case 'chunk':
                  if (parsed.content) {
                    onChunk(parsed.content)
                  }
                  break
                case 'citations':
                  if (parsed.citations && onCitations) {
                    onCitations(parsed.citations as Citation[])
                  }
                  break
                case 'error':
                  if (parsed.message && onError) {
                    onError(parsed.message)
                  }
                  break
                case 'done':
                  onDone?.()
                  return
              }
            } catch {
              // Ignore parse errors for non-JSON data
            }
          }
        }
      }
    } finally {
      reader.releaseLock()
    }
  }

  return { sendMessage, sessionId }
}
