import { defineStore } from 'pinia'
import { ref, shallowRef } from 'vue'
import * as api from '@/api/index.js'
import { useSessionsStore } from '@/stores/sessions.js'

// Guard against late completion of stale loadSession calls
let _lastLoadId = null

async function _tryGenerateName(sessionId) {
  try {
    const sessionsStore = useSessionsStore()
    const session = sessionsStore.sessions.find(s => s.session_id === sessionId)
    if (session?.name) return  // already has a name
    const { name } = await api.generateSessionName(sessionId)
    if (name) sessionsStore.updateName(sessionId, name)
  } catch {
    // fire-and-forget — ignore errors
  }
}

export const useChatStore = defineStore('chat', () => {
  const currentId = ref(null)
  const currentProfileId = ref(null)
  const messages = ref([])
  const streaming = ref(false)
  const pendingImages = ref([])   // base64 strings
  const pendingFiles = ref([])    // { name, path, size, content_type }
  const contextTokens = ref(0)
  const maxContextTokens = ref(0)
  const loading = ref(false)

  // The in-progress streaming message (direct ref for perf)
  const streamingMsg = shallowRef(null)
  // True while replaying buffered events from a reconnect — suppresses animations
  const replayMode = ref(false)

  async function loadSession(id) {
    if (currentId.value === id) return
    _lastLoadId = id
    loading.value = true
    messages.value = []
    streaming.value = false
    streamingMsg.value = null
    contextTokens.value = 0
    maxContextTokens.value = 0

    try {
      const session = await api.getSession(id)
      // Ignore stale responses — user may have switched to a different session while this loaded
      if (_lastLoadId !== id) return
      currentProfileId.value = session.profile_id ?? null
      messages.value = buildMessageList(session.messages ?? [])
      await mergeFeedback(id)
      if (session.context_token_count) contextTokens.value = session.context_token_count
      if (session.max_context_tokens)  maxContextTokens.value = session.max_context_tokens
      // Set currentId and hash AFTER REST completes — this triggers ws.connect() via the
      // watch in ChatArea. Doing it last guarantees messages are populated before WS replay
      // starts, eliminating the race where WS replay arrives before REST and gets overwritten.
      currentId.value = id
      location.hash = id
    } finally {
      loading.value = false
    }
  }

  function clearSession() {
    currentId.value = null
    currentProfileId.value = null
    messages.value = []
    streaming.value = false
    streamingMsg.value = null
    contextTokens.value = 0
    maxContextTokens.value = 0
    location.hash = ''
  }

  // Force-reload session history from the server, bypassing the same-id guard.
  // Called after WebSocket reconnect so the client sees the saved response
  // even if streaming finished while the client was disconnected.
  async function reloadSession(id) {
    if (!id) return
    try {
      const session = await api.getSession(id)
      currentProfileId.value = session.profile_id ?? null
      messages.value = buildMessageList(session.messages ?? [])
      await mergeFeedback(id)
      if (session.context_token_count) contextTokens.value = session.context_token_count
      if (session.max_context_tokens) maxContextTokens.value = session.max_context_tokens
      // Clear any orphaned streaming message so subsequent deltas don't mutate a ghost object
      streamingMsg.value = null
      streaming.value = false
    } catch {
      // Silently ignore — stale data is better than a crash
    }
  }

  // Pull stored ratings for this session and stamp them onto built messages.
  // Each assistant block has id `h_<index>` where index is its first raw position;
  // that index is also the feedback key (session_id, message_index).
  async function mergeFeedback(id) {
    try {
      const { feedback = [] } = await api.getFeedback(id)
      if (!feedback.length) return
      const map = new Map(feedback.map(f => [f.message_index, f.rating]))
      for (const m of messages.value) {
        if (m.role !== 'assistant') continue
        const idx = m.id?.startsWith?.('h_') ? Number(m.id.slice(2)) : NaN
        if (Number.isInteger(idx) && map.has(idx)) {
          m.rating = map.get(idx)
        }
      }
    } catch {
      // Feedback is non-critical — ignore failures so chat still works.
    }
  }

  // Toggle / set feedback on an assistant block. Optimistic — local state
  // updates first, server call follows; rolls back on failure.
  async function rateMessage(msg, rating) {
    if (!currentId.value || !msg) return
    const idx = msg.id?.startsWith?.('h_') ? Number(msg.id.slice(2)) : NaN
    if (!Number.isInteger(idx)) return
    // Toggle off if user clicks the active rating again
    const next = msg.rating === rating ? 0 : rating
    const prev = msg.rating ?? 0
    msg.rating = next
    try {
      await api.setFeedback(currentId.value, idx, next)
    } catch (err) {
      msg.rating = prev
      throw err
    }
  }

  function saveDraft(text) {
    if (!currentId.value) return
    if (text) localStorage.setItem(`draft:${currentId.value}`, text)
    else localStorage.removeItem(`draft:${currentId.value}`)
  }

  function loadDraft(id) {
    return localStorage.getItem(`draft:${id}`) ?? ''
  }

  // ─── WS event handlers ──────────────────────────────────────────────────

  function onStreamStart() {
    // On reconnect: remove the incomplete frozen message so replay rebuilds it cleanly.
    // On a fresh start streamingMsg is null, so this branch is skipped.
    if (streamingMsg.value) {
      const idx = messages.value.indexOf(streamingMsg.value)
      if (idx !== -1) messages.value.splice(idx, 1)
      streamingMsg.value = null
    }
    streaming.value = true
    const msg = {
      id: `stream_${Date.now()}`,
      role: 'assistant',
      type: 'stream',
      thinking: null,    // { text, done } | null
      tools: [],         // tool card objects
      text: '',
      done: false,
      time: new Date().toISOString(),
      animate: !replayMode.value,  // no cursor animation during replay
      statusLabel: null,
    }
    messages.value.push(msg)
    // Store the reactive proxy (from the array) so mutations trigger Vue updates
    streamingMsg.value = messages.value[messages.value.length - 1]
  }

  function onReplayStart() {
    replayMode.value = true
  }

  function onReplayEnd() {
    replayMode.value = false
    // Restore animate on the message now that live events will follow
    if (streamingMsg.value) streamingMsg.value.animate = true
  }

  function onThinkingDelta(delta) {
    const msg = streamingMsg.value
    if (!msg) return
    if (!msg.thinking) msg.thinking = { text: '', done: false }
    msg.thinking.text += delta
  }

  function onThinkingEnd() {
    const msg = streamingMsg.value
    if (msg?.thinking) msg.thinking.done = true
  }

  // Find the last spawn_agent card that's still collecting subagent events
  function _lastSpawnCard(msg) {
    // Custom findLast for older browser compatibility (ES2023 polyfill-free)
    for (let i = msg.tools.length - 1; i >= 0; i--) {
      const t = msg.tools[i]
      if (t.kind === 'tool' && t.name === 'spawn_agent') return t
    }
    return undefined
  }

  function onTurnThinking(data) {
    const msg = streamingMsg.value
    if (!msg) return
    // Shape item so ThinkingCard can receive it directly as :msg (stable reference, no new object each render)
    const item = { kind: 'turn_thinking', isSubagent: data.is_subagent ?? false, thinking: { text: data.thinking ?? '', done: true } }
    if (data.is_subagent) {
      const spawn = _lastSpawnCard(msg)
      if (spawn) { spawn.steps.push(item); return }
    }
    msg.tools.push(item)
  }

  function onPlanningStatus(data) {
    const msg = streamingMsg.value
    if (!msg) return
    if (data.is_subagent) {
      const spawn = _lastSpawnCard(msg)
      if (spawn) spawn.planningLabel = data.label ?? ''
      return  // never bleed subagent planning into parent UI
    }
    msg.statusLabel = data.label ?? ''
  }

  function onPlanReady(data) {
    const msg = streamingMsg.value
    if (!msg) return
    if (data.is_subagent) {
      const spawn = _lastSpawnCard(msg)
      if (spawn) {
        spawn.planningLabel = null
        spawn.steps.push({ kind: 'plan', text: data.plan ?? '' })
      }
      return  // never bleed subagent plan into parent UI
    }
    msg.statusLabel = null
    msg.tools.push({ kind: 'plan', text: data.plan ?? '' })
  }

  function onToolStarted(data) {
    const msg = streamingMsg.value
    if (!msg) return
    const card = {
      kind: 'tool',
      id: `tool_${Date.now()}`,
      name: data.tool,
      args: data.args,
      result: null,
      success: null,
      pending: true,
      isSubagent: data.is_subagent ?? false,
      steps: []
    }
    if (data.is_subagent) {
      const spawn = _lastSpawnCard(msg)
      if (spawn) { spawn.steps.push(card); return }
    }
    msg.tools.push(card)
  }

  function onToolCall(data) {
    const msg = streamingMsg.value
    if (!msg) return
    // Sub-agent tool call → update inside spawn_agent steps
    if (data.is_subagent) {
      const spawn = _lastSpawnCard(msg)
      if (spawn) {
        let step = null
        for (let i = spawn.steps.length - 1; i >= 0; i--) {
          const t = spawn.steps[i]
          if (t.kind === 'tool' && t.name === data.tool && t.pending) { step = t; break }
        }
        if (step) {
          step.result = data.result
          step.success = data.success !== false
          step.pending = false
          return
        }
      }
    }
    // Top-level tool call
    let card = null
    for (let i = msg.tools.length - 1; i >= 0; i--) {
      const t = msg.tools[i]
      if (t.kind === 'tool' && t.name === data.tool && t.pending) { card = t; break }
    }
    if (card) {
      card.result = data.result
      card.success = data.success !== false
      card.pending = false
      if (data.metadata) card.metadata = data.metadata
    }
  }

  function onStreamDelta(delta) {
    const msg = streamingMsg.value
    if (!msg) return
    if (msg.statusLabel) msg.statusLabel = null
    msg.text += delta
  }

  function onStreamEnd(data) {
    const msg = streamingMsg.value
    if (msg) {
      msg.done = true
      msg.elapsed_seconds = data?.elapsed_seconds ?? null
      msg.tool_call_count  = data?.tool_call_count  ?? null
      msg.token_count      = data?.token_count      ?? null
      streamingMsg.value = null
      // Purge if nothing was ever written (avoids empty msg-assistant divs)
      if (!msg.thinking && !msg.tools.length && !msg.text) {
        messages.value = messages.value.filter(m => m !== msg)
      }
    }
    streaming.value = false

    if (data?.context_tokens) contextTokens.value = data.context_tokens
    if (data?.max_context_tokens) maxContextTokens.value = data.max_context_tokens

    // Update session preview
    if (currentId.value && msg?.text) {
      useSessionsStore().updatePreview(currentId.value, msg.text.slice(0, 80))
    }

    // Try to generate session name in background (only if not named yet)
    if (currentId.value) {
      _tryGenerateName(currentId.value)
    }
  }

  function onStreamStopped() {
    const msg = streamingMsg.value
    if (msg) {
      msg.done = true
      streamingMsg.value = null
      if (!msg.thinking && !msg.tools.length && !msg.text) {
        messages.value = messages.value.filter(m => m !== msg)
      }
    }
    streaming.value = false
  }

  function onProfileSwitched(data) {
    currentProfileId.value = data.profile_id
  }

  function onContextCompressed(data) {
    messages.value.push({
      id: `compress_${Date.now()}`,
      role: 'system',
      type: 'compression_notice',
      before: data.messages_before,
      after: data.messages_after,
      summary: data.summary ?? ''
    })
  }

  function onError(data) {
    streaming.value = false
    streamingMsg.value = null
    messages.value.push({
      id: `err_${Date.now()}`,
      role: 'system',
      type: 'error',
      text: data.message ?? 'An error occurred'
    })
  }

  function appendUserMessage(text, images, files) {
    messages.value.push({
      id: `user_${Date.now()}`,
      role: 'user',
      text,
      images: [...images],
      files: [...files],
      time: new Date().toISOString(),
      animate: true
    })
  }

  // ─── Helpers ────────────────────────────────────────────────────────────

  // The server returns a flat array:
  //   { role:'assistant', tool_calls:[{id,name,arguments}] }  ← tool request (no content)
  //   { role:'tool', tool_call_id, name, content }            ← tool result
  //   { role:'assistant', content }                           ← final text response
  //
  // We group them into single AssistantMessage objects with tools[] + text.
  function buildMessageList(raw) {
    const result = []
    let i = 0

    while (i < raw.length) {
      const m = raw[i]

      // Compression notice
      if (m.is_compression) {
        result.push({ id: `h_${i}`, role: 'system', type: 'compression_notice', summary: m.content ?? '' })
        i++; continue
      }

      // Skip bare system messages
      if (m.role === 'system') { i++; continue }

      if (m.is_summary) {
        result.push({ id: `h_${i}`, role: 'assistant', type: 'summary', text: m.content ?? '' })
        i++; continue
      }

      if (m.role === 'user') {
        const imgs = (m.images ?? []).map(b =>
          (typeof b === 'string' && b.startsWith('data:')) ? b : `data:image/jpeg;base64,${String(b ?? '')}`
        )
        result.push({ id: `h_${i}`, role: 'user', text: m.content ?? '', images: imgs, files: m.files ?? [], time: m.created_at ?? null })
        i++; continue
      }

      if (m.role === 'assistant') {
        // Collect all tool-call rounds + final content into one message
        const msgId = `h_${i}`
        const tools = []
        let thinking = null
        let text = ''
        let time = null

        while (i < raw.length && raw[i].role === 'assistant') {
          const am = raw[i]

          // Plan card — stored as a separate is_plan message, inject into tools array
          if (am.is_plan) {
            tools.push({ kind: 'plan', text: am.content ?? '' })
            i++
            continue
          }

          // Accumulate thinking from any assistant turn (first one wins for display)
          if (am.thinking && !thinking) {
            thinking = { text: am.thinking, done: true }
          }

          if (am.tool_calls?.length) {
            // Build a lookup map for this round's tool calls
            const callMap = {}
            for (const tc of am.tool_calls) {
              const card = {
                kind: 'tool',
                id: tc.id,
                name: tc.name,
                args: tc.arguments ?? {},
                result: null,
                success: true,
                pending: false,
                isSubagent: false,
                steps: []
              }
              tools.push(card)
              callMap[tc.id] = card
            }
            i++
            // Collect matching tool results
            while (i < raw.length && raw[i].role === 'tool') {
              const tr = raw[i]
              const card = callMap[tr.tool_call_id]
              if (card) {
                card.result = tr.content ?? ''
                card.success = !tr.content?.startsWith('Error:')
                card.metadata = tr.metadata ?? {}
              }
              i++
            }
          } else {
            const content = am.content ?? ''
            if (!time && am.created_at) time = am.created_at
            i++
            if (content) {
              // Non-empty text → this is the final response
              text = content
              break
            }
            // Empty content: if no more assistant messages follow, we're done
            if (i >= raw.length || raw[i].role !== 'assistant') {
              break
            }
            // Otherwise skip this empty intermediate message and keep accumulating
          }
        }

        // Only add the message if there's something to show
        if (thinking || tools.length || text) {
          // Scan the group for metrics (saved on the final assistant turn)
          let elapsed_seconds = null, tool_call_count = null, token_count = null
          const startIdx = Number(msgId.slice(2))
          for (let j = startIdx; j < i; j++) {
            const am = raw[j]
            if (!am || am.role !== 'assistant') continue
            if (am.elapsed_seconds != null) elapsed_seconds = am.elapsed_seconds
            if (am.tool_call_count != null) tool_call_count = am.tool_call_count
            if (am.token_count     != null) token_count     = am.token_count
          }
          result.push({ id: msgId, role: 'assistant', type: 'history', thinking, tools, text, done: true, time, elapsed_seconds, tool_call_count, token_count })
        }
        continue
      }

      // Orphan tool message (shouldn't happen)
      i++
    }

    return result
  }

  return {
    currentId,
    currentProfileId,
    messages,
    streaming,
    pendingImages,
    pendingFiles,
    contextTokens,
    maxContextTokens,
    loading,
    streamingMsg,
    loadSession,
    clearSession,
    reloadSession,
    rateMessage,
    saveDraft,
    loadDraft,
    appendUserMessage,
    onStreamStart,
    onReplayStart,
    onReplayEnd,
    onThinkingDelta,
    onThinkingEnd,
    onTurnThinking,
    onPlanningStatus,
    onPlanReady,
    onToolStarted,
    onToolCall,
    onStreamDelta,
    onStreamEnd,
    onStreamStopped,
    onProfileSwitched,
    onContextCompressed,
    onError
  }
})