import type { Message, ToolCall } from '#cli/core/model'
import { isEditLikeToolName, isReadLikeToolName } from '#cli/core/tool-metadata'
import {
  INITIAL_EXPANDED_TOOL_OUTPUT_LINES,
  LOAD_MORE_TOOL_OUTPUT_LINES,
  clampExpandedToolOutputLines,
} from '../tool-output-pagination'

type AssistantMessage = Extract<Message, { role: 'assistant' }>
type ToolCallPart = Extract<AssistantMessage['parts'][number], { type: 'tool_call' }>

interface ToolUiState {
  messages: Message[]
  selectedToolCallId: string | null
  expandedBashOutputs: Set<string>
  expandedToolOutputLineCounts: Map<string, number>
  unifiedEditDiffToolCallIds: Set<string>
  collapsedToolCallIds: Set<string>
  expandedToolCallGroups: Set<string>
}

type UpdateAssistant = (
  assistantId: string,
  updater: (m: AssistantMessage) => AssistantMessage,
) => void

export function createToolUiActions(
  opts: Readonly<{
    state: ToolUiState
    updateAssistant: UpdateAssistant
    scheduleToolCallReadHighlightByCallId: (opts: {
      state: Pick<
        ToolUiState,
        'messages' | 'selectedToolCallId' | 'expandedBashOutputs'
      >
      callId: string
      mode: 'collapsed' | 'expanded'
      updateAssistant: UpdateAssistant
    }) => void
    scheduleToolCallDiffHighlightByCallId: (opts: {
      state: Pick<ToolUiState, 'messages'>
      callId: string
      updateAssistant: UpdateAssistant
    }) => void
  }>,
): Readonly<{
  selectToolCall: (callId: string | null) => void
  toggleBashOutput: (toolCallId: string) => void
  loadMoreBashOutput: (toolCallId: string) => void
  toggleEditDiffLayout: (toolCallId: string) => void
  toggleToolCallCollapsed: (toolCallId: string) => void
  toggleToolCallGroupCollapsed: (groupKey: string) => void
  toggleToolCallGroup: (groupKey: string) => void
  clearCollapsedToolCalls: () => void
  setAllThinkingCollapsed: (collapsed: boolean) => void
}> {
  const state = opts.state

  function findToolCallPart(
    message: AssistantMessage,
    callId: string,
  ): ToolCallPart | null {
    return (
      message.parts.find(
        (part): part is ToolCallPart =>
          part.type === 'tool_call' && part.call.id === callId,
      ) ?? null
    )
  }

  function findToolCallInTranscript(
    callId: string,
  ): { assistantId: string; call: ToolCall } | null {
    const id = String(callId ?? '').trim()
    if (!id) return null
    for (const message of state.messages) {
      if (message.role !== 'assistant') continue
      const part = findToolCallPart(message, id)
      if (part) return { assistantId: message.id, call: part.call }
    }
    return null
  }

  function selectToolCall(callId: string | null): void {
    state.selectedToolCallId = callId
    if (callId) {
      opts.scheduleToolCallReadHighlightByCallId({
        state,
        callId,
        mode: 'collapsed',
        updateAssistant: opts.updateAssistant,
      })
      opts.scheduleToolCallDiffHighlightByCallId({
        state,
        callId,
        updateAssistant: opts.updateAssistant,
      })
    }
  }

  function toggleBashOutput(toolCallId: string): void {
    const found = findToolCallInTranscript(toolCallId)
    const toolNameLower = String(found?.call?.name ?? '')
      .trim()
      .toLowerCase()
    // Read results are already truncated upstream before they reach the CLI transcript.
    // Keep the current non-expandable UX until the raw-output pipeline changes.
    if (isReadLikeToolName(toolNameLower)) return

    const newSet = new Set(state.expandedBashOutputs)
    const nextLineCounts = new Map(state.expandedToolOutputLineCounts)
    if (newSet.has(toolCallId)) {
      newSet.delete(toolCallId)
      nextLineCounts.delete(toolCallId)
    } else {
      newSet.add(toolCallId)
      nextLineCounts.set(
        toolCallId,
        clampExpandedToolOutputLines(INITIAL_EXPANDED_TOOL_OUTPUT_LINES),
      )
    }
    state.expandedBashOutputs = newSet
    state.expandedToolOutputLineCounts = nextLineCounts
    if (newSet.has(toolCallId)) {
      opts.scheduleToolCallReadHighlightByCallId({
        state,
        callId: toolCallId,
        mode: 'expanded',
        updateAssistant: opts.updateAssistant,
      })
    }
  }

  function loadMoreBashOutput(toolCallId: string): void {
    const id = String(toolCallId ?? '').trim()
    if (!id || !state.expandedBashOutputs.has(id)) return
    const next = new Map(state.expandedToolOutputLineCounts)
    const current = clampExpandedToolOutputLines(
      Number(next.get(id) ?? INITIAL_EXPANDED_TOOL_OUTPUT_LINES),
    )
    next.set(
      id,
      clampExpandedToolOutputLines(current + LOAD_MORE_TOOL_OUTPUT_LINES),
    )
    state.expandedToolOutputLineCounts = next

    const found = findToolCallInTranscript(id)
    const toolNameLower = String(found?.call?.name ?? '')
      .trim()
      .toLowerCase()
    if (isReadLikeToolName(toolNameLower)) {
      opts.scheduleToolCallReadHighlightByCallId({
        state,
        callId: id,
        mode: 'expanded',
        updateAssistant: opts.updateAssistant,
      })
    }
  }

  function toggleEditDiffLayout(toolCallId: string): void {
    const found = findToolCallInTranscript(toolCallId)
    const toolNameLower = String(found?.call?.name ?? '')
      .trim()
      .toLowerCase()
    if (!isEditLikeToolName(toolNameLower)) return

    const next = new Set(state.unifiedEditDiffToolCallIds)
    if (next.has(toolCallId)) next.delete(toolCallId)
    else next.add(toolCallId)
    state.unifiedEditDiffToolCallIds = next
  }

  function toggleToolCallCollapsed(toolCallId: string): void {
    const id = String(toolCallId ?? '').trim()
    if (!id) return
    const next = new Set(state.collapsedToolCallIds)
    if (next.has(id)) next.delete(id)
    else next.add(id)
    state.collapsedToolCallIds = next
  }

  /**
   * Toggle collapse for a tool-call group header in expand mode.
   *
   * Because `groupKey === child1Id`, we can't use the same bit for both
   * "group is collapsed" and "child1 is collapsed". Instead:
   * - Collapsing: add groupKey + all childIds to `collapsedToolCallIds`;
   *   remove groupKey from `expandedToolCallGroups`.
   * - Expanding: add groupKey to `expandedToolCallGroups` (acts as an
   *   override that signals "group is open"); leave `collapsedToolCallIds`
   *   untouched so all children remain collapsed when the group reopens.
   *
   * `isGroupCollapsed` is updated to check both sets:
   *   collapsed = collapsedToolCallIds.has(key) && !expandedToolCallGroups.has(key)
   */
  function toggleToolCallGroupCollapsed(groupKey: string): void {
    const id = String(groupKey ?? '').trim()
    if (!id) return

    // The group is currently expanded when it's NOT in collapsedToolCallIds,
    // OR when expandedToolCallGroups overrides a previous collapse.
    const groupIsExpanded =
      !state.collapsedToolCallIds.has(id) ||
      state.expandedToolCallGroups.has(id)

    if (groupIsExpanded) {
      // Collapsing: mark group + all children as collapsed.
      const nextCollapsed = new Set(state.collapsedToolCallIds)
      nextCollapsed.add(id) // child1Id === groupKey
      const childIds = findGroupChildCallIds(id) // child2, child3, …
      for (const cid of childIds) nextCollapsed.add(cid)
      state.collapsedToolCallIds = nextCollapsed

      // Remove the explicit "expanded" override if present.
      if (state.expandedToolCallGroups.has(id)) {
        const nextExpanded = new Set(state.expandedToolCallGroups)
        nextExpanded.delete(id)
        state.expandedToolCallGroups = nextExpanded
      }
    } else {
      // Expanding: set the override flag. Children stay in collapsedToolCallIds
      // so they appear collapsed when the group re-opens.
      const nextExpanded = new Set(state.expandedToolCallGroups)
      nextExpanded.add(id)
      state.expandedToolCallGroups = nextExpanded
    }
  }

  /**
   * Find all callIds that belong to the same consecutive tool-call group
   * as `groupKey` (the first callId in the group), excluding `groupKey` itself.
   */
  function findGroupChildCallIds(groupKey: string): string[] {
    for (const message of state.messages) {
      if (message.role !== 'assistant') continue
      const parts = message.parts
      let foundIdx = -1
      for (let i = 0; i < parts.length; i++) {
        const part = parts[i]
        if (part.type === 'tool_call' && part.call.id === groupKey) {
          foundIdx = i
          break
        }
      }
      if (foundIdx < 0) continue
      const groupPart = parts[foundIdx]
      if (groupPart.type !== 'tool_call') continue
      const toolName = String(groupPart.call.name ?? '').trim().toLowerCase()
      const childIds: string[] = []
      for (let i = foundIdx + 1; i < parts.length; i++) {
        const p = parts[i]
        if (p.type === 'tool_call') {
          if (String(p.call.name ?? '').trim().toLowerCase() === toolName) {
            childIds.push(p.call.id)
          } else {
            break
          }
        } else if (
          p.type === 'tool_output'
          || p.type === 'tool_result'
          || p.type === 'approve'
        ) {
          continue
        } else {
          break
        }
      }
      return childIds
    }
    return []
  }

  function toggleToolCallGroup(groupKey: string): void {
    const id = String(groupKey ?? '').trim()
    if (!id) return
    const next = new Set(state.expandedToolCallGroups)
    if (next.has(id)) {
      next.delete(id)
      // In normal mode, `isGroupCollapsed` auto-expands the group when
      // `selectedToolCallId` belongs to it. Clear the selection so the
      // explicit collapse actually takes effect.
      if (state.selectedToolCallId) {
        const childIds = findGroupChildCallIds(id)
        if (
          state.selectedToolCallId === id ||
          childIds.includes(state.selectedToolCallId)
        ) {
          state.selectedToolCallId = null
        }
      }
    } else {
      next.add(id)
    }
    state.expandedToolCallGroups = next
  }

  function clearCollapsedToolCalls(): void {
    if (state.collapsedToolCallIds.size === 0) return
    state.collapsedToolCallIds = new Set()
  }

  function setAllThinkingCollapsed(collapsed: boolean): void {
    for (const msg of state.messages) {
      if (msg.role !== 'assistant') continue
      opts.updateAssistant(msg.id, (m) => {
        let changed = false
        const parts = m.parts.map((p) => {
          if (p.type !== 'thinking') return p
          if (p.collapsed === collapsed) return p
          changed = true
          return { ...p, collapsed }
        })
        return changed ? { ...m, parts } : m
      })
    }
  }

  return {
    selectToolCall,
    toggleBashOutput,
    loadMoreBashOutput,
    toggleEditDiffLayout,
    toggleToolCallCollapsed,
    toggleToolCallGroupCollapsed,
    toggleToolCallGroup,
    clearCollapsedToolCalls,
    setAllThinkingCollapsed,
  }
}
