import type { DisplayMessage, ToolCallInfo, RenderedUIPanel } from '../types/index'

/**
 * Convert SDK Message[] into flat DisplayMessage[] for UI.
 *
 * Critical invariants:
 *
 * 1. **Index parity with the source array**: `display[i]` corresponds to
 *    `messages[i]`. Callers (useAgent) attach `_toolCalls` / `_uiPanels` /
 *    `_attachments` from the raw message by index, so we MUST NOT skip
 *    messages — empty turns become empty-text DisplayMessages (which the UI
 *    renders as blank but preserved so tool-call metadata stays aligned).
 *
 * 2. **toolUse and toolResult are first-class**: previously we flattened
 *    toolUse blocks into ugly `*[tool: name]*` placeholder text, which
 *    destroyed the input payload and broke pairing with toolResult. Now we
 *    extract them into a proper `toolCalls: ToolCallInfo[]` on the
 *    DisplayMessage, so ToolCallBubble can show full input/output.
 *
 * 3. **render_ui is re-materialized from toolUse input**: the `render_ui`
 *    tool's input IS the HTML/CSS/JS of the panel. On re-hydration from IDB
 *    (new tab, page reload), `_uiPanels` may be absent but the toolUse
 *    blocks survive — so we rebuild `uiPanels` from them. That's how the
 *    UI panels survive forever, not just within the live stream.
 *
 * 4. **Text extraction preserves ordering**: toolUse/toolResult blocks are
 *    NOT inlined as text. They become siblings of the message via
 *    toolCalls/uiPanels. Text remains text.
 */
export function toDisplayMessages(messages: any[]): DisplayMessage[] {
  const result: DisplayMessage[] = []

  // First pass: build preliminary list + collect tool-result-by-id lookup
  // so we can pair toolUse (in one assistant msg) with toolResult (in the
  // next user msg, per SDK convention).
  const toolResultsById = new Map<string, { status: 'success' | 'error'; text: string }>()
  for (const msg of messages) {
    const blocks = Array.isArray(msg.content) ? msg.content : []
    for (const b of blocks) {
      if (b.toolResult?.toolUseId) {
        const txt = Array.isArray(b.toolResult.content)
          ? b.toolResult.content.map((c: any) => c.text || (typeof c === 'string' ? c : JSON.stringify(c))).join('')
          : String(b.toolResult.content ?? '')
        const status = b.toolResult.status === 'error' ? 'error' : 'success'
        toolResultsById.set(b.toolResult.toolUseId, { status, text: txt })
      }
    }
  }

  for (let i = 0; i < messages.length; i++) {
    const msg = messages[i]
    const role = msg.role === 'user' ? 'user' : 'assistant'
    const blocks = Array.isArray(msg.content) ? msg.content : [{ text: String(msg.content ?? '') }]

    const texts: string[] = []
    const toolCalls: ToolCallInfo[] = []
    const inlinePanels: RenderedUIPanel[] = []

    for (const b of blocks) {
      if (b.text) {
        texts.push(b.text)
      } else if (b.toolUse) {
        const id = b.toolUse.toolUseId || `tu-${i}-${toolCalls.length}`
        const paired = toolResultsById.get(b.toolUse.toolUseId || '')
        toolCalls.push({
          id,
          toolUseId: b.toolUse.toolUseId || id,
          name: b.toolUse.name,
          input: b.toolUse.input ?? {},
          status: paired ? paired.status : 'success',
          result: paired?.text,
          startedAt: 0,
          endedAt: 0,
        })

        // Re-materialize render_ui panels from persisted tool input so they
        // survive tab reload / IDB roundtrip / context compaction.
        if (b.toolUse.name === 'render_ui' && b.toolUse.input) {
          const inp = b.toolUse.input
          inlinePanels.push({
            id: `panel-${id}`,
            title: inp.title || 'panel',
            html: inp.html || '',
            css: inp.css,
            js: inp.js,
            createdAt: 0,
          })
        }
      }
      // toolResult blocks are consumed via the lookup map above; we don't
      // render them as their own DisplayMessage text.
    }

    const text = texts.join('').trim()
    const entry: DisplayMessage = { id: `msg-${i}`, role, text }

    // Prefer explicit metadata (stashed during live streaming) but fall back
    // to what we extracted from raw blocks for IDB-reloaded messages.
    const rawToolCalls = (msg as any)._toolCalls as ToolCallInfo[] | undefined
    const rawUiPanels = (msg as any)._uiPanels as RenderedUIPanel[] | undefined
    const rawAttachments = (msg as any)._attachments

    if (rawToolCalls?.length) entry.toolCalls = rawToolCalls
    else if (toolCalls.length) entry.toolCalls = toolCalls

    if (rawUiPanels?.length) entry.uiPanels = rawUiPanels
    else if (inlinePanels.length) entry.uiPanels = inlinePanels

    if (rawAttachments?.length) entry.attachments = rawAttachments

    result.push(entry)
  }
  return result
}

/**
 * Serialize agent.messages safely for IndexedDB.
 *
 * Preserves:
 * - text blocks verbatim
 * - toolUse blocks with full input payload (so render_ui can be rehydrated)
 * - toolResult blocks with toolUseId + status + content
 * - image/document blocks (binary → base64)
 * - our hack metadata: _toolCalls, _uiPanels, _attachments
 */
export function serializeMessages(messages: any[]): any[] {
  return messages.map(msg => ({
    role: msg.role,
    ...(msg._toolCalls ? { _toolCalls: msg._toolCalls } : {}),
    ...(msg._uiPanels ? { _uiPanels: msg._uiPanels } : {}),
    ...(msg._attachments ? { _attachments: msg._attachments } : {}),
    content: Array.isArray(msg.content) ? msg.content.map((b: any) => {
      if (typeof b.toJSON === 'function') return b.toJSON()
      if (b.text !== undefined) return { text: b.text }
      if (b.toolUse) return { toolUse: { name: b.toolUse.name, toolUseId: b.toolUse.toolUseId, input: b.toolUse.input } }
      if (b.toolResult) return {
        toolResult: {
          toolUseId: b.toolResult.toolUseId,
          status: b.toolResult.status,
          content: b.toolResult.content?.map((c: any) => ({ text: c.text || String(c) })) || [],
        },
      }
      if (b.image) {
        let bytes = b.image.source?.bytes
        if (bytes instanceof Uint8Array) {
          let s = ''
          for (let i = 0; i < bytes.length; i++) s += String.fromCharCode(bytes[i])
          bytes = btoa(s)
        }
        return { image: { format: b.image.format, source: { bytes } } }
      }
      if (b.document) {
        let bytes = b.document.source?.bytes
        if (bytes instanceof Uint8Array) {
          let s = ''
          for (let i = 0; i < bytes.length; i++) s += String.fromCharCode(bytes[i])
          bytes = btoa(s)
        }
        return { document: { name: b.document.name, format: b.document.format, source: { bytes } } }
      }
      return { text: String(b.text ?? '') }
    }) : [{ text: String(msg.content ?? '') }]
  }))
}