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

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)

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

    // Update URL hash
    location.hash = id

    try {
      const session = await api.getSession(id)
      currentProfileId.value = session.profile_id ?? null
      messages.value = buildMessageList(session.messages ?? [])

      // Restore draft
      const draft = localStorage.getItem(`draft:${id}`) ?? ''
      _draftRestoreCallback?.(draft)
    } finally {
      loading.value = false
    }
  }

  // Draft restore callback — set by InputBar
  let _draftRestoreCallback = null
  function onDraftRestore(cb) { _draftRestoreCallback = cb }

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

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

  function onStreamStart() {
    // Close any previous streaming message (prevents double cursor on reconnect)
    if (streamingMsg.value) {
      streamingMsg.value.done = true
      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
    }
    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 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) {
    return [...msg.tools].reverse().find(t => t.kind === 'tool' && t.name === 'spawn_agent')
  }

  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 onPlanReady(data) {
    const msg = streamingMsg.value
    if (!msg) return
    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) {
        const step = spawn.steps.findLast(t => t.kind === 'tool' && t.name === data.tool && t.pending)
        if (step) {
          step.result = data.result
          step.success = data.success !== false
          step.pending = false
          return
        }
      }
    }
    // Top-level tool call
    const card = msg.tools.findLast(t => t.kind === 'tool' && t.name === data.tool && t.pending)
    if (card) {
      card.result = data.result
      card.success = data.success !== false
      card.pending = false
    }
  }

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

  function onStreamEnd(data) {
    const msg = streamingMsg.value
    if (msg) {
      msg.done = true
      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))
    }
  }

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

  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]
    })
  }

  // ─── 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]

      // 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 =>
          b.startsWith('data:') ? b : `data:image/jpeg;base64,${b}`
        )
        result.push({ id: `h_${i}`, role: 'user', text: m.content ?? '', images: imgs, files: m.files ?? [] })
        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 = ''

        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:')
              }
              i++
            }
          } else {
            const content = am.content ?? ''
            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) {
          result.push({ id: msgId, role: 'assistant', type: 'history', thinking, tools, text, done: true })
        }
        continue
      }

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

    return result
  }

  return {
    currentId,
    currentProfileId,
    messages,
    streaming,
    pendingImages,
    pendingFiles,
    contextTokens,
    maxContextTokens,
    loading,
    streamingMsg,
    loadSession,
    saveDraft,
    onDraftRestore,
    appendUserMessage,
    onStreamStart,
    onThinkingDelta,
    onThinkingEnd,
    onTurnThinking,
    onPlanReady,
    onToolStarted,
    onToolCall,
    onStreamDelta,
    onStreamEnd,
    onStreamStopped,
    onProfileSwitched,
    onContextCompressed,
    onError
  }
})