import { defineStore } from 'pinia'
import { ref } from 'vue'
import api from '../api'

export interface TaskInfo {
  task_id: string
  name: string
  status: string
  progress: number
  message: string
  result_files: string[]
  error: string | null
  log_lines: string[]
}

export const useProcessingStore = defineStore('processing', () => {
  const currentTask = ref<TaskInfo | null>(null)
  const tasks = ref<TaskInfo[]>([])
  const logs = ref<string[]>([])
  const taskSource = ref<string>('')

  let ws: WebSocket | null = null
  let reconnectAttempts = 0
  let reconnectTimer: ReturnType<typeof setTimeout> | null = null
  let currentTaskId: string | null = null
  const MAX_RECONNECT = 5

  function connectWebSocket(taskId: string) {
    disconnectWebSocket()
    currentTaskId = taskId
    reconnectAttempts = 0
    doConnect(taskId)
  }

  function doConnect(taskId: string) {
    const token = localStorage.getItem('token')
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    const host = window.location.host
    const url = `${protocol}//${host}/ws/task/${taskId}?token=${token}`

    ws = new WebSocket(url)

    ws.onopen = () => {
      reconnectAttempts = 0
    }

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data)
        if (data.error) return // ignore error messages from ws
        currentTask.value = data
        logs.value = data.log_lines || []

        const idx = tasks.value.findIndex(t => t.task_id === data.task_id)
        if (idx >= 0) {
          tasks.value[idx] = data
        } else {
          tasks.value.unshift(data)
        }

        if (data.status === 'completed' || data.status === 'failed') {
          setTimeout(() => disconnectWebSocket(), 2000)
        }
      } catch {}
    }

    ws.onerror = () => {
      // Error will be followed by onclose, which handles reconnection
    }

    ws.onclose = (event) => {
      ws = null
      // Auto-reconnect if task is still running and not manually disconnected
      const task = currentTask.value
      if (
        currentTaskId === taskId &&
        task &&
        (task.status === 'pending' || task.status === 'running') &&
        reconnectAttempts < MAX_RECONNECT
      ) {
        const delay = Math.min(1000 * Math.pow(2, reconnectAttempts), 10000)
        reconnectAttempts++
        reconnectTimer = setTimeout(() => {
          if (currentTaskId === taskId) {
            doConnect(taskId)
          }
        }, delay)
      }
    }
  }

  function disconnectWebSocket() {
    currentTaskId = null
    reconnectAttempts = MAX_RECONNECT // prevent reconnect
    if (reconnectTimer) {
      clearTimeout(reconnectTimer)
      reconnectTimer = null
    }
    if (ws) {
      ws.close()
      ws = null
    }
  }

  async function startTask(endpoint: string, body?: any, source: string = 'processing') {
    const res = await api.post(endpoint, body || {})
    const taskId = res.data.task_id
    taskSource.value = source
    currentTask.value = {
      task_id: taskId,
      name: res.data.message || '',
      status: 'pending',
      progress: 0,
      message: '',
      result_files: [],
      error: null,
      log_lines: [],
    }
    logs.value = []
    connectWebSocket(taskId)
    return taskId
  }

  async function pollTaskStatus(taskId: string) {
    const res = await api.get(`/processing/status/${taskId}`)
    return res.data
  }

  return { currentTask, tasks, logs, taskSource, connectWebSocket, disconnectWebSocket, startTask, pollTaskStatus }
})
