export interface WsPlanUpdate {
  type: 'PlanUpdate'
  workflow_id: string
  steps: { order: number; description: string; command: string; status?: string }[]
}

export interface WsStepStatusUpdate {
  type: 'StepStatusUpdate'
  step_id: string
  status: string
  output: string
}

export interface WsWorkflowStatusUpdate {
  type: 'WorkflowStatusUpdate'
  workflow_id: string
  status: string
}

export interface WsRequirementUpdate {
  type: 'RequirementUpdate'
  workflow_id: string
  requirement: string
}

export interface WsReportReady {
  type: 'ReportReady'
  workflow_id: string
}

export interface WsProjectUpdate {
  type: 'ProjectUpdate'
  project_id: string
  name: string
}

export interface WsLlmCallLog {
  type: 'LlmCallLog'
  workflow_id: string
  entry: import('./types').LlmCallLogEntry
}

export interface WsActivityUpdate {
  type: 'ActivityUpdate'
  workflow_id: string
  activity: string
}

export interface WsError {
  type: 'Error'
  message: string
}

export type WsMessage = WsPlanUpdate | WsStepStatusUpdate | WsWorkflowStatusUpdate | WsRequirementUpdate | WsReportReady | WsProjectUpdate | WsLlmCallLog | WsActivityUpdate | WsError

export type WsHandler = (msg: WsMessage) => void

export function connectWs(projectId: string, onMessage: WsHandler): { close: () => void } {
  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:'
  const wsBase = import.meta.env.BASE_URL.replace(/\/$/, '')
  const url = `${proto}//${location.host}/ws${wsBase}/${projectId}`
  let ws: WebSocket | null = null
  let reconnectTimer: ReturnType<typeof setTimeout> | null = null
  let closed = false

  function connect() {
    if (closed) return
    ws = new WebSocket(url)

    ws.onmessage = (e) => {
      try {
        const msg: WsMessage = JSON.parse(e.data)
        onMessage(msg)
      } catch {
        // ignore malformed messages
      }
    }

    ws.onclose = () => {
      if (!closed) {
        reconnectTimer = setTimeout(connect, 2000)
      }
    }

    ws.onerror = () => {
      ws?.close()
    }
  }

  connect()

  return {
    close() {
      closed = true
      if (reconnectTimer) clearTimeout(reconnectTimer)
      ws?.close()
    },
  }
}