polylan-submitter/polylan_submitter/src/composables/useSubmissionsWebSocket.ts

import { ref, computed, onBeforeUnmount } from 'vue'
import { useWebSocket } from '@vueuse/core'

interface SubmissionUpdate {
  id: string
  scope_id?: string
  notes?: string
  files?: SubmissionFileUpdate[]
  created_at: string
  updated_at: string
}

interface SubmissionFileUpdate {
  id: string
  original_filename: string
  content_type: string
  file_size: number
  file_url?: string
  ocr_status: string
  ocr_raw_data: string
  crops?: ImageCropUpdate[]
}

interface ImageCropUpdate {
  id: string
  x1: number
  y1: number
  x2: number
  y2: number
  label: string
  ocr_status: string
  ocr_raw_data: string
  created_at: string
}

type WebSocketData = SubmissionUpdate | SubmissionFileUpdate | ImageCropUpdate

export function useSubmissionsWebSocket(userSubmissions: any) {
  const isConnected = ref(false)
  const reconnectAttempts = ref(0)
  const maxReconnectAttempts = 10
  const reconnectDelay = ref(3000)

  const wsUrl = computed(() => {
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    return `${protocol}//${window.location.host}/ws/submissions/`
  })

  const {
    status,
    open,
    close,
  } = useWebSocket(wsUrl.value, {
    onConnected: () => {
      isConnected.value = true
      reconnectAttempts.value = 0
      reconnectDelay.value = 3000
      console.log('✓ Connected to submissions WebSocket')
    },
    onDisconnected: () => {
      isConnected.value = false
      console.log('✗ Disconnected from submissions WebSocket, attempting to reconnect...')
      attemptReconnect()
    },
    onError: (event) => {
      console.error('WebSocket error:', event)
    },
    onMessage: (_ws, event) => {
      try {
        const messageData = event.data as string
        const data = JSON.parse(messageData) as WebSocketData

        if ('files' in data) {
          handleSubmissionUpdate(data as SubmissionUpdate)
        } else if ('original_filename' in data) {
          handleFileUpdate(data as SubmissionFileUpdate)
        } else if ('x1' in data) {
          handleCropUpdate(data as ImageCropUpdate)
        }
      } catch (error) {
        console.error('Failed to parse WebSocket message:', error)
      }
    },
    immediate: true,
  })

  const attemptReconnect = () => {
    if (reconnectAttempts.value >= maxReconnectAttempts) {
      console.error('Max reconnection attempts reached')
      return
    }

    reconnectAttempts.value++
    const delay = reconnectDelay.value * Math.pow(1.5, reconnectAttempts.value - 1)

    console.log(`Reconnecting (attempt ${reconnectAttempts.value}/${maxReconnectAttempts}) in ${Math.round(delay / 1000)}s...`)

    setTimeout(() => {
      open()
    }, delay)
  }

  const handleSubmissionUpdate = (submission: SubmissionUpdate) => {
    const existingIndex = userSubmissions.value.findIndex((s: any) => s.id === submission.id)

    if (existingIndex >= 0) {
      // Merge into existing object to preserve references
      Object.assign(userSubmissions.value[existingIndex], submission)
    } else {
      userSubmissions.value.unshift(submission)
    }

    console.log('Updated submission:', submission.id)
  }

  const handleFileUpdate = (file: SubmissionFileUpdate) => {
    userSubmissions.value.forEach((submission: any) => {
      if (!submission.files) return

      const fileIndex = submission.files.findIndex((f: any) => f.id === file.id)
      if (fileIndex >= 0) {
        // Merge into existing object to preserve references
        Object.assign(submission.files[fileIndex], file)
        console.log('Updated file:', file.id)
      }
    })
  }

  const handleCropUpdate = (crop: ImageCropUpdate) => {
    userSubmissions.value.forEach((submission: any) => {
      if (!submission.files) return

      submission.files.forEach((file: any) => {
        if (!file.crops) return

        const cropIndex = file.crops.findIndex((c: any) => c.id === crop.id)
        if (cropIndex >= 0) {
          // Merge into existing object to preserve references
          Object.assign(file.crops[cropIndex], crop)
          console.log('Updated crop:', crop.id)
        }
      })
    })
  }

  onBeforeUnmount(() => {
    close()
  })

  return {
    isConnected,
    status,
    open,
    close,
  }
}