import { defineStore } from 'pinia'
import { ref, computed } 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[]
}

interface TaskConnection {
  ws: WebSocket | null
  reconnectAttempts: number
  reconnectTimer: ReturnType<typeof setTimeout> | null
}

export const useProcessingStore = defineStore('processing', () => {
  // --- Multi-task tracking ---
  const activeTasks = ref(new Map<string, TaskInfo>())

  const activeTaskList = computed(() =>
    Array.from(activeTasks.value.values())
  )

  const currentTask = computed<TaskInfo | null>(() =>
    activeTaskList.value[0] ?? null
  )

  // --- Legacy compatibility ---
  const tasks = ref<TaskInfo[]>([])
  const logs = ref<string[]>([])
  const taskSource = ref<string>('')

  // --- Per-task WebSocket management ---
  const taskConnections = new Map<string, TaskConnection>()
  const MAX_RECONNECT = 5

  function connectWebSocket(taskId: string) {
    disconnectTaskWS(taskId)
    taskConnections.set(taskId, { ws: null, reconnectAttempts: 0, reconnectTimer: null })
    doConnect(taskId)
  }

  function doConnect(taskId: string) {
    const conn = taskConnections.get(taskId)
    if (!conn) return

    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}`

    const socket = new WebSocket(url)
    conn.ws = socket

    socket.onopen = () => {
      conn.reconnectAttempts = 0
    }

    socket.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data)
        if (data.error) return

        // Update activeTasks map
        activeTasks.value.set(data.task_id, data)

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

        // Legacy: update logs for the current (most recent) task
        if (currentTask.value?.task_id === data.task_id) {
          logs.value = data.log_lines || []
        }

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

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

    socket.onclose = () => {
      conn.ws = null
      const task = activeTasks.value.get(taskId)
      if (
        task &&
        (task.status === 'pending' || task.status === 'running') &&
        conn.reconnectAttempts < MAX_RECONNECT
      ) {
        const delay = Math.min(1000 * Math.pow(2, conn.reconnectAttempts), 10000)
        conn.reconnectAttempts++
        conn.reconnectTimer = setTimeout(() => {
          if (taskConnections.has(taskId)) {
            doConnect(taskId)
          }
        }, delay)
      }
    }
  }

  function disconnectTaskWS(taskId: string) {
    const conn = taskConnections.get(taskId)
    if (!conn) return
    conn.reconnectAttempts = MAX_RECONNECT // prevent reconnect
    if (conn.reconnectTimer) {
      clearTimeout(conn.reconnectTimer)
      conn.reconnectTimer = null
    }
    if (conn.ws) {
      conn.ws.close()
      conn.ws = null
    }
    taskConnections.delete(taskId)
  }

  /** Disconnect all task WebSockets (backward compat) */
  function disconnectWebSocket() {
    for (const taskId of Array.from(taskConnections.keys())) {
      disconnectTaskWS(taskId)
    }
  }

  function removeTask(taskId: string) {
    disconnectTaskWS(taskId)
    activeTasks.value.delete(taskId)
    const idx = tasks.value.findIndex(t => t.task_id === taskId)
    if (idx >= 0) tasks.value.splice(idx, 1)
  }

  async function retryTask(taskId: string) {
    const res = await api.post(`/api/tasks/${taskId}/retry`)
    const newTaskId: string = res.data.task_id
    const taskInfo: TaskInfo = {
      task_id: newTaskId,
      name: res.data.message || '',
      status: 'pending',
      progress: 0,
      message: '',
      result_files: [],
      error: null,
      log_lines: [],
    }
    activeTasks.value.set(newTaskId, taskInfo)
    connectWebSocket(newTaskId)
    return newTaskId
  }

  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
    const taskInfo: TaskInfo = {
      task_id: taskId,
      name: res.data.message || '',
      status: 'pending',
      progress: 0,
      message: '',
      result_files: [],
      error: null,
      log_lines: [],
    }
    activeTasks.value.set(taskId, taskInfo)
    logs.value = []
    connectWebSocket(taskId)
    return taskId
  }

  return {
    activeTasks,
    activeTaskList,
    currentTask,
    tasks,
    logs,
    taskSource,
    connectWebSocket,
    disconnectWebSocket,
    startTask,
    removeTask,
    retryTask,
  }
})
