import { ref, computed, watch } from 'vue'
import type { Ref } from 'vue'
import type { DataSource } from 'mpr-ui-framework'
import { usePipelineStore } from '../stores/pipeline'
import { useStageRegistry } from './useStageRegistry'

interface CheckpointFrame {
  seq: number
  timestamp: number
  jpeg_b64: string
}

export function useCheckpointLoader(
  jobId: Ref<string>,
  source: DataSource,
) {
  const pipeline = usePipelineStore()

  const currentFrameImage = ref<string | null>(null)
  const currentFrameRef = ref<number | null>(null)

  const checkpointFrames = ref<CheckpointFrame[]>([])
  const checkpointFrameIndex = ref(0)
  const checkpointStage = ref<string | null>(null)

  const stripSelStart = ref(0)
  const stripSelEndOverride = ref<number | null>(null)
  const stripSelEnd = computed(() =>
    stripSelEndOverride.value ?? Math.max(0, checkpointFrames.value.length - 1),
  )

  // Track current frame from SSE
  source.on<{ frame_ref: number; jpeg_b64: string }>('frame_update', (e) => {
    currentFrameImage.value = e.jpeg_b64
    currentFrameRef.value = e.frame_ref
  })

  async function loadCheckpoint(job: string, stage: string) {
    try {
      const lookupId = pipeline.timelineId || job

      const resp = await fetch(`/api/detect/checkpoints/${lookupId}/${stage}`)
      if (!resp.ok) return

      const data = await resp.json()
      checkpointFrames.value = data.frames ?? []
      checkpointStage.value = stage

      if (checkpointFrames.value.length > 0) {
        checkpointFrameIndex.value = 0
        const first = checkpointFrames.value[0]
        currentFrameImage.value = first.jpeg_b64
        currentFrameRef.value = first.seq
      }

      stripSelStart.value = 0
      stripSelEndOverride.value = null
    } catch (e) {
      console.error('Failed to load checkpoint:', e)
    }
  }

  function setCheckpointFrame(index: number) {
    if (index < 0 || index >= checkpointFrames.value.length) return
    checkpointFrameIndex.value = index
    const frame = checkpointFrames.value[index]
    currentFrameImage.value = frame.jpeg_b64
    currentFrameRef.value = frame.seq
  }

  const { stages, checkpointStageFor } = useStageRegistry()

  // Auto-load checkpoint when entering editor mode.
  watch(
    () => [pipeline.layoutMode, pipeline.editorStage, jobId.value, stages.value.length] as const,
    ([mode, stage, job]) => {
      if (mode === 'bbox_editor' && stage && job) {
        const cpStage = checkpointStageFor(stage)
        if (cpStage) {
          loadCheckpoint(job, cpStage)
        }
      }
    },
    { immediate: true },
  )

  return {
    currentFrameImage,
    currentFrameRef,
    checkpointFrames,
    checkpointFrameIndex,
    checkpointStage,
    stripSelStart,
    stripSelEnd,
    stripSelEndOverride,
    loadCheckpoint,
    setCheckpointFrame,
  }
}