PPT/src/views/Editor/EnglishSpeaking/composables/useDialogueEngine.ts

import { ref, reactive, computed, onUnmounted } from 'vue'
import type { PreviewChatMessage, DialogueAPI, SessionConfig, DialogueReport } from '@/types/englishSpeaking'
import { MockDialogueAPI, RealDialogueAPI } from '../services/llmService'

export function useDialogueEngine(mode: 'preview' | 'real' = 'preview') {
  const messages = ref<PreviewChatMessage[]>([])
  const sessionId = ref<string | null>(null)
  const expiresAt = ref<string | null>(null)
  const currentRound = ref(1)
  const isComplete = ref(false)
  const countdownSeconds = ref<number | null>(null)

  let api: DialogueAPI = mode === 'real' ? new RealDialogueAPI() : new MockDialogueAPI()
  let currentAbortController: AbortController | null = null
  let countdownTimer: ReturnType<typeof setInterval> | null = null
  let ttsUtterance: SpeechSynthesisUtterance | null = null

  const isProcessing = computed(() => messages.value.some(m => m.status === 'loading'))
  const canRecord = computed(() => !isProcessing.value && !isComplete.value)

  // ==================== Session ====================

  async function initSession(config: SessionConfig) {
    try {
      const info = await api.createSession(config)
      sessionId.value = info.sessionId
      expiresAt.value = info.expiresAt || null

      messages.value.push({
        id: crypto.randomUUID(),
        role: 'ai',
        content: info.aiMessage,
        timestamp: new Date(),
        status: 'done',
      })

      if (info.expiresAt) startCountdown(info.expiresAt)
      speakTTS(info.aiMessage)
    } catch (err: any) {
      console.error('Failed to init session:', err)
    }
  }

  // ==================== Send Message ====================

  async function sendStudentMessage(audioBlob: Blob) {
    if (!sessionId.value || isProcessing.value) return

    // Add student message (loading)
    const studentMsg = reactive<PreviewChatMessage>({
      id: crypto.randomUUID(),
      role: 'student',
      content: '',
      timestamp: new Date(),
      status: 'loading',
      audioBlob,
    })
    messages.value.push(studentMsg)

    // Add AI message placeholder
    const aiMsg = reactive<PreviewChatMessage>({
      id: crypto.randomUUID(),
      role: 'ai',
      content: '',
      timestamp: new Date(),
      status: 'loading',
    })

    currentAbortController = new AbortController()

    try {
      const generator = api.speak(sessionId.value, audioBlob, currentAbortController.signal)

      for await (const event of generator) {
        if (event.type === 'transcript') {
          studentMsg.content = event.text
          studentMsg.status = 'done'
          // Now push AI message placeholder
          messages.value.push(aiMsg)
        } else if (event.type === 'token') {
          aiMsg.content += event.text
        } else if (event.type === 'done') {
          aiMsg.status = 'done'
          isComplete.value = event.isComplete

          if (!event.isComplete) {
            currentRound.value++
          }

          speakTTS(aiMsg.content)
        }
      }

      // If student message never got transcript, mark done with fallback
      if (studentMsg.status === 'loading') {
        studentMsg.status = 'done'
      }
      if (aiMsg.status === 'loading') {
        aiMsg.status = 'done'
      }
    } catch (err: any) {
      if (err.name === 'AbortError') return

      // Determine which message to mark as error
      if (studentMsg.status === 'loading') {
        studentMsg.status = 'error'
        studentMsg.error = err.message || 'Request failed'
      } else if (aiMsg.status === 'loading') {
        aiMsg.status = 'error'
        aiMsg.error = err.message || 'Request failed'
      }
    } finally {
      currentAbortController = null
    }
  }

  // ==================== Retry / Regenerate ====================

  async function retryMessage(messageId: string) {
    const msg = messages.value.find(m => m.id === messageId)
    if (!msg || msg.status !== 'error') return

    if (msg.role === 'student' && msg.audioBlob) {
      // Remove the failed student message and any subsequent AI message
      const idx = messages.value.indexOf(msg)
      messages.value.splice(idx)
      await sendStudentMessage(msg.audioBlob)
    }
  }

  async function regenerateAiMessage(messageId: string) {
    const msg = messages.value.find(m => m.id === messageId)
    if (!msg || msg.role !== 'ai' || msg.status !== 'error') return

    // Find the student message before this AI message
    const idx = messages.value.indexOf(msg)
    const prevStudent = messages.value.slice(0, idx).reverse().find(m => m.role === 'student')
    if (!prevStudent?.audioBlob || !sessionId.value) return

    // Remove the failed AI message
    messages.value.splice(idx, 1)

    // Re-add AI placeholder and stream
    const aiMsg = reactive<PreviewChatMessage>({
      id: crypto.randomUUID(),
      role: 'ai',
      content: '',
      timestamp: new Date(),
      status: 'loading',
    })
    messages.value.push(aiMsg)

    currentAbortController = new AbortController()

    try {
      const generator = api.speak(sessionId.value, prevStudent.audioBlob, currentAbortController.signal)

      for await (const event of generator) {
        if (event.type === 'transcript') {
          // Skip transcript on regenerate, student message already exists
        } else if (event.type === 'token') {
          aiMsg.content += event.text
        } else if (event.type === 'done') {
          aiMsg.status = 'done'
          isComplete.value = event.isComplete
          if (!event.isComplete) currentRound.value++
          speakTTS(aiMsg.content)
        }
      }

      if (aiMsg.status === 'loading') aiMsg.status = 'done'
    } catch (err: any) {
      if (err.name === 'AbortError') return
      aiMsg.status = 'error'
      aiMsg.error = err.message || 'Request failed'
    } finally {
      currentAbortController = null
    }
  }

  // ==================== Report ====================

  function getReport(): Promise<DialogueReport> {
    if (!sessionId.value) return Promise.reject(new Error('No session'))

    return new Promise((resolve, reject) => {
      let attempts = 0
      const maxAttempts = 15 // 30s / 2s

      const poll = async () => {
        attempts++
        try {
          const report = await api.getReport(sessionId.value!)
          resolve(report)
        } catch {
          if (attempts >= maxAttempts) {
            reject(new Error('Report timeout'))
          } else {
            setTimeout(poll, 2000)
          }
        }
      }
      poll()
    })
  }

  // ==================== TTS ====================

  function speakTTS(text: string) {
    if (!text || typeof speechSynthesis === 'undefined') return

    cancelTTS()
    ttsUtterance = new SpeechSynthesisUtterance(text)
    ttsUtterance.lang = 'en-US'
    ttsUtterance.rate = 0.9
    speechSynthesis.speak(ttsUtterance)
  }

  function cancelTTS() {
    if (typeof speechSynthesis !== 'undefined') {
      speechSynthesis.cancel()
    }
    ttsUtterance = null
  }

  // ==================== Countdown ====================

  function startCountdown(expiresAtStr: string) {
    stopCountdown()

    const update = () => {
      const remaining = Math.max(0, Math.floor((new Date(expiresAtStr).getTime() - Date.now()) / 1000))
      countdownSeconds.value = remaining
      if (remaining <= 0) {
        stopCountdown()
        isComplete.value = true
      }
    }

    update()
    countdownTimer = setInterval(update, 1000)
  }

  function stopCountdown() {
    if (countdownTimer) { clearInterval(countdownTimer); countdownTimer = null }
    countdownSeconds.value = null
  }

  // ==================== Abort ====================

  function abort() {
    currentAbortController?.abort()
    currentAbortController = null
  }

  // ==================== Streaming Speak (WebSocket) ====================

  /** 流式开始：立即 push 学生占位消息，返回 controller 供外部推 chunk / 结束 */
  function beginStudentStream(opts: {
    sampleRate: number
    bits?: number
    channels?: number
  }) {
    if (!sessionId.value || isProcessing.value) {
      return null
    }

    // 立即占位：录音按完成的那一刻 UI 就已经显示学生泡泡 + AI placeholder
    const studentMsg = reactive<PreviewChatMessage>({
      id: crypto.randomUUID(),
      role: 'student',
      content: '',
      timestamp: new Date(),
      status: 'loading',
    })
    messages.value.push(studentMsg)

    const aiMsg = reactive<PreviewChatMessage>({
      id: crypto.randomUUID(),
      role: 'ai',
      content: '',
      timestamp: new Date(),
      status: 'loading',
    })
    messages.value.push(aiMsg)

    const wsUrl = buildWsUrl('/speak-stream')
    const ws = new WebSocket(wsUrl)
    ws.binaryType = 'arraybuffer'

    let aborted = false
    let chunkQueue: ArrayBuffer[] = []
    let open = false

    const finalizeError = (msg: string) => {
      if (studentMsg.status === 'loading') {
        studentMsg.status = 'error'
        studentMsg.error = msg
      }
      else if (aiMsg.status === 'loading') {
        aiMsg.status = 'error'
        aiMsg.error = msg
      }
    }

    ws.onopen = () => {
      open = true
      ws.send(JSON.stringify({
        type: 'start',
        sessionId: sessionId.value,
        sampleRate: opts.sampleRate,
        bits: opts.bits ?? 16,
        channels: opts.channels ?? 1,
      }))
      // flush 队列里攒的 chunk
      for (const c of chunkQueue) ws.send(c)
      chunkQueue = []
    }

    ws.onmessage = (e: MessageEvent) => {
      try {
        const data = JSON.parse(e.data)
        if (data.type === 'transcript') {
          studentMsg.content = data.text
          studentMsg.status = 'done'
        }
        else if (data.type === 'token') {
          aiMsg.content += data.content
        }
        else if (data.type === 'done') {
          aiMsg.status = 'done'
          isComplete.value = !!data.isComplete
          if (!data.isComplete) currentRound.value++
          speakTTS(aiMsg.content)
          ws.close()
        }
        else if (data.type === 'error') {
          finalizeError(friendlyErrorMessage(data.message))
          ws.close()
        }
      } catch { /* ignore */ }
    }

    ws.onerror = () => {
      if (!aborted) finalizeError(friendlyErrorMessage('WebSocket error'))
    }
    ws.onclose = () => {
      if (studentMsg.status === 'loading') finalizeError(friendlyErrorMessage('Connection closed'))
      else if (aiMsg.status === 'loading') finalizeError(friendlyErrorMessage('Connection closed'))
    }

    const pushChunk = (chunk: ArrayBuffer) => {
      if (aborted) return
      if (open && ws.readyState === WebSocket.OPEN) ws.send(chunk)
      else chunkQueue.push(chunk)
    }

    const finish = () => {
      if (aborted) return
      if (open && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ type: 'stop' }))
      } else {
        // 还没 open 就被要求结束 → 废话不说 直接 close
        ws.close()
      }
    }

    const abortStream = () => {
      aborted = true
      try { ws.close() } catch { /* ignore */ }
      finalizeError('Aborted')
    }

    // 便于重试：保存关联的学生消息 id
    return { studentMsgId: studentMsg.id, aiMsgId: aiMsg.id, pushChunk, finish, abortStream }
  }

  /**
   * 流式失败时的 HTTP fallback：用完整 audioBlob 走旧 /speak 路径。
   * 会把 beginStudentStream 已 push 的占位消息回收（避免重复）。
   */
  async function streamFallback(audioBlob: Blob, studentMsgId: string, aiMsgId: string) {
    // 移除占位消息
    messages.value = messages.value.filter(m => m.id !== studentMsgId && m.id !== aiMsgId)
    // 走旧流程
    await sendStudentMessage(audioBlob)
  }

  // ==================== Cleanup ====================

  onUnmounted(() => {
    abort()
    cancelTTS()
    stopCountdown()
  })

  return {
    messages,
    sessionId,
    currentRound,
    isComplete,
    isProcessing,
    canRecord,
    countdownSeconds,

    initSession,
    sendStudentMessage,
    beginStudentStream,
    streamFallback,
    retryMessage,
    regenerateAiMessage,
    getReport,
    abort,
    cancelTTS,
  }
}

// ==================== Helpers ====================

/** 把 /api/speaking/dialogue/... 的 HTTP base URL 转成 ws/wss URL */
function buildWsUrl(path: string): string {
  const API_BASE = 'http://localhost:8000/api/speaking/dialogue'
  const wsBase = API_BASE.replace(/^http/, 'ws')
  return wsBase + path
}

/** 把后端英文错误转成面向用户的中文友好文案 */
function friendlyErrorMessage(raw: string | undefined): string {
  const map: Record<string, string> = {
    'No speech detected': '没听清，请再说一次',
    'Session not found': '会话已失效，请刷新页面重新开始',
    'Session is not active': '会话已结束',
    'sessionId required': '会话参数缺失',
    'First message must be start': '协议异常，请刷新重试',
    'Invalid start payload': '协议异常，请刷新重试',
    'Internal error': '服务暂时不可用，请稍后重试',
    'WebSocket error': '网络连接异常',
    'Connection closed': '连接中断，请重试',
    'Aborted': '已取消',
  }
  if (!raw) return '请求失败，请重试'
  return map[raw] || raw
}