import { batch } from "solid-js"
import { createStore, produce, reconcile } from "solid-js/store"
import type { SetStoreFunction } from "solid-js/store"
import { getLogger } from "../../lib/logger"
import type { ClientPart, MessageInfo } from "../../types/message"
import { clearRecordDisplayCacheForMessages } from "./record-display-cache"
import type {
  InstanceMessageState,
  LatestTodoSnapshot,
  MessageRecord,
  MessageUpsertInput,
  PartUpdateInput,
  PendingPartEntry,
  PermissionEntry,
  QuestionEntry,
  ReplaceMessageIdOptions,
  ScrollSnapshot,
  SessionRecord,
  SessionUpsertInput,
  SessionUsageState,
  UsageEntry,
} from "./types"

const storeLog = getLogger("session")

interface MessageStoreHooks {
  onSessionCleared?: (instanceId: string, sessionId: string) => void
}

function createInitialState(instanceId: string): InstanceMessageState {
  return {
    instanceId,
    sessions: {},
    sessionOrder: [],
    messages: {},
    lastAssistantMessageIds: {},
    messageInfoVersion: {},
    pendingParts: {},
    sessionRevisions: {},
    permissions: {
      queue: [],
      active: null,
      byMessage: {},
    },
    questions: {
      queue: [],
      active: null,
      byMessage: {},
    },
    usage: {},
    scrollState: {},
    latestTodos: {},
  }
}

function ensurePartId(messageId: string, part: ClientPart, index: number): string {
  if (typeof part.id === "string" && part.id.length > 0) {
    return part.id
  }

  if (part.type === "tool") {
    throw new Error("Tool part missing id")
  }

  const fallbackId = `${messageId}-part-${index}`
  part.id = fallbackId
  return fallbackId
}

const PENDING_PART_MAX_AGE_MS = 30_000

function clonePart(part: ClientPart): ClientPart {
  // Cloning is intentionally disabled; message parts
  // are stored as received from the backend.
  return part
}

function cloneStructuredValue<T>(value: T): T {
  // Legacy helper kept as a no-op to avoid deep copies.
  return value
}

function areMessageIdListsEqual(a: string[], b: string[]): boolean {
  if (a.length !== b.length) {
    return false
  }
  for (let index = 0; index < a.length; index++) {
    if (a[index] !== b[index]) {
      return false
    }
  }
  return true
}

function createEmptyUsageState(): SessionUsageState {
  return {
    entries: {},
    totalInputTokens: 0,
    totalOutputTokens: 0,
    totalReasoningTokens: 0,
    totalCost: 0,
    actualUsageTokens: 0,
    latestMessageId: undefined,
  }
}

function extractUsageEntry(info: MessageInfo | undefined): UsageEntry | null {
  if (!info || info.role !== "assistant") return null
  const messageId = typeof info.id === "string" ? info.id : undefined
  if (!messageId) return null
  const tokens = info.tokens
  if (!tokens) return null
  const inputTokens = tokens.input ?? 0
  const outputTokens = tokens.output ?? 0
  const reasoningTokens = tokens.reasoning ?? 0
  const cacheReadTokens = tokens.cache?.read ?? 0
  const cacheWriteTokens = tokens.cache?.write ?? 0
  if (inputTokens === 0 && outputTokens === 0 && reasoningTokens === 0 && cacheReadTokens === 0 && cacheWriteTokens === 0) {
    return null
  }
  const combinedTokens = info.summary ? outputTokens : inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens
  return {
    messageId,
    inputTokens,
    outputTokens,
    reasoningTokens,
    cacheReadTokens,
    cacheWriteTokens,
    combinedTokens,
    cost: info.cost ?? 0,
    timestamp: info.time?.created ?? 0,
    hasContextUsage: inputTokens + cacheReadTokens + cacheWriteTokens > 0,
  }
}

function applyUsageState(state: SessionUsageState, entry: UsageEntry | null) {
  if (!entry) return
  state.entries[entry.messageId] = entry
  state.totalInputTokens += entry.inputTokens
  state.totalOutputTokens += entry.outputTokens
  state.totalReasoningTokens += entry.reasoningTokens
  state.totalCost += entry.cost
  if (!state.latestMessageId || entry.timestamp >= (state.entries[state.latestMessageId]?.timestamp ?? 0)) {
    state.latestMessageId = entry.messageId
    state.actualUsageTokens = entry.combinedTokens
  }
}

function removeUsageEntry(state: SessionUsageState, messageId: string | undefined) {
  if (!messageId) return
  const existing = state.entries[messageId]
  if (!existing) return
  state.totalInputTokens -= existing.inputTokens
  state.totalOutputTokens -= existing.outputTokens
  state.totalReasoningTokens -= existing.reasoningTokens
  state.totalCost -= existing.cost
  delete state.entries[messageId]
  if (state.latestMessageId === messageId) {
    state.latestMessageId = undefined
    state.actualUsageTokens = 0
    let latest: UsageEntry | null = null
    for (const candidate of Object.values(state.entries) as UsageEntry[]) {
      if (!latest || candidate.timestamp >= latest.timestamp) {
        latest = candidate
      }
    }
    if (latest) {
      state.latestMessageId = latest.messageId
      state.actualUsageTokens = latest.combinedTokens
    }
  }
}

function rebuildUsageStateFromInfos(infos: Iterable<MessageInfo>): SessionUsageState {
  const usageState = createEmptyUsageState()
  for (const info of infos) {
    const entry = extractUsageEntry(info)
    if (entry) {
      applyUsageState(usageState, entry)
    }
  }
  return usageState
}

export interface InstanceMessageStore {
  instanceId: string
  state: InstanceMessageState
  setState: SetStoreFunction<InstanceMessageState>
  addOrUpdateSession: (input: SessionUpsertInput) => void
  hydrateMessages: (sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable<MessageInfo>) => void
  upsertMessage: (input: MessageUpsertInput) => void
  applyPartUpdate: (input: PartUpdateInput) => void
  applyPartDelta: (input: {
    messageId: string
    partId: string
    field: string
    delta: string
    bumpRevision?: boolean
    bumpSessionRevision: boolean
  }) => void
  removeMessage: (messageId: string) => void
  removeMessagePart: (messageId: string, partId: string) => void
  bufferPendingPart: (entry: PendingPartEntry) => void
  flushPendingParts: (messageId: string) => void
  replaceMessageId: (options: ReplaceMessageIdOptions) => void
  setMessageInfo: (messageId: string, info: MessageInfo) => void
  getMessageInfo: (messageId: string) => MessageInfo | undefined
  upsertPermission: (entry: PermissionEntry) => void
  removePermission: (permissionId: string) => void
  getPermissionState: (messageId?: string, partId?: string) => { entry: PermissionEntry; active: boolean } | null
  upsertQuestion: (entry: QuestionEntry) => void
  removeQuestion: (requestId: string) => void
  getQuestionState: (messageId?: string, partId?: string) => { entry: QuestionEntry; active: boolean } | null
  setSessionRevert: (sessionId: string, revert?: SessionRecord["revert"] | null) => void
  getSessionRevert: (sessionId: string) => SessionRecord["revert"] | undefined | null
  rebuildUsage: (sessionId: string, infos: Iterable<MessageInfo>) => void
  getSessionUsage: (sessionId: string) => SessionUsageState | undefined
  setScrollSnapshot: (sessionId: string, scope: string, snapshot: Omit<ScrollSnapshot, "updatedAt">) => void
  getScrollSnapshot: (sessionId: string, scope: string) => ScrollSnapshot | undefined
  getSessionRevision: (sessionId: string) => number
  getSessionMessageIds: (sessionId: string) => string[]
  getLastAssistantMessageId: (sessionId: string) => string | undefined
  // Index of the most recent message in the session that contains a compaction part.
  // Returns -1 if there has been no compaction.
  getLastCompactionMessageIndex: (sessionId: string) => number
  getMessage: (messageId: string) => MessageRecord | undefined
  getLatestTodoSnapshot: (sessionId: string) => LatestTodoSnapshot | undefined
  clearSession: (sessionId: string) => void
  clearInstance: () => void
}

export function createInstanceMessageStore(instanceId: string, hooks?: MessageStoreHooks): InstanceMessageStore {
  const [state, setState] = createStore<InstanceMessageState>(createInitialState(instanceId))

  const TODO_TOOL_NAME = "todowrite"

  const messageInfoCache = new Map<string, MessageInfo>()

  function findLastAssistantMessageId(messageIds: readonly string[]): string | undefined {
    for (let index = messageIds.length - 1; index >= 0; index -= 1) {
      const messageId = messageIds[index]
      if (state.messages[messageId]?.role === "assistant") {
        return messageId
      }
    }
    return undefined
  }

  function recomputeLastAssistantMessageId(sessionId: string, messageIds?: readonly string[]) {
    if (!sessionId) return
    setState("lastAssistantMessageIds", sessionId, findLastAssistantMessageId(messageIds ?? state.sessions[sessionId]?.messageIds ?? []))
  }

  function getLastCompactionMessageIndex(sessionId: string): number {
    if (!sessionId) return -1
    const ids = state.sessions[sessionId]?.messageIds ?? []
    // Scan from the end: we only care about the most recent compaction.
    for (let i = ids.length - 1; i >= 0; i--) {
      const messageId = ids[i]
      const record = state.messages[messageId]
      if (!record || !Array.isArray(record.partIds) || record.partIds.length === 0) continue
      for (const partId of record.partIds) {
        const part = record.parts[partId]?.data
        if ((part as any)?.type === "compaction") {
          return i
        }
      }
    }
    return -1
  }

  function isCompletedTodoPart(part: ClientPart | undefined): boolean {
    if (!part || (part as any).type !== "tool") {
      return false
    }
    const toolName = typeof (part as any).tool === "string" ? (part as any).tool : ""
    if (toolName !== TODO_TOOL_NAME) {
      return false
    }
    const toolState = (part as any).state
    if (!toolState || typeof toolState !== "object") {
      return false
    }
    return (toolState as { status?: string }).status === "completed"
  }

  function recordLatestTodoSnapshot(sessionId: string, snapshot: LatestTodoSnapshot) {
    if (!sessionId) return
    setState("latestTodos", sessionId, (existing) => {
      if (existing && existing.timestamp > snapshot.timestamp) {
        return existing
      }
      return snapshot
    })
  }

  function maybeUpdateLatestTodoFromRecord(record: MessageRecord | undefined) {
    if (!record || !Array.isArray(record.partIds) || record.partIds.length === 0) {
      return
    }
    for (let index = record.partIds.length - 1; index >= 0; index -= 1) {
      const partId = record.partIds[index]
      const partRecord = record.parts[partId]
      if (!partRecord) continue
      if (isCompletedTodoPart(partRecord.data)) {
        const timestamp = typeof record.updatedAt === "number" ? record.updatedAt : Date.now()
        recordLatestTodoSnapshot(record.sessionId, { messageId: record.id, partId, timestamp })
        break
      }
    }
  }

  function clearLatestTodoSnapshot(sessionId: string) {
    setState("latestTodos", sessionId, undefined)
  }

  function bumpSessionRevision(sessionId: string) {
    if (!sessionId) return
    setState("sessionRevisions", sessionId, (value = 0) => value + 1)
  }

  function getSessionRevisionValue(sessionId: string) {
    return state.sessionRevisions[sessionId] ?? 0
  }

  function getLastAssistantMessageIdValue(sessionId: string) {
    return state.lastAssistantMessageIds[sessionId]
  }

  function withUsageState(sessionId: string, updater: (draft: SessionUsageState) => void) {
    setState("usage", sessionId, (current) => {
      const draft = current
        ? {
            ...current,
            entries: { ...current.entries },
          }
        : createEmptyUsageState()
      updater(draft)
      return draft
    })
  }

  function updateUsageWithInfo(info: MessageInfo | undefined) {
    if (!info || typeof info.sessionID !== "string") return
    const messageId = typeof info.id === "string" ? info.id : undefined
    if (!messageId) return
    withUsageState(info.sessionID, (draft) => {
      removeUsageEntry(draft, messageId)
      const entry = extractUsageEntry(info)
      if (entry) {
        applyUsageState(draft, entry)
      }
    })
  }

  function rebuildUsage(sessionId: string, infos: Iterable<MessageInfo>) {
    const usageState = rebuildUsageStateFromInfos(infos)
    setState("usage", sessionId, usageState)
  }

  function getSessionUsage(sessionId: string) {
    return state.usage[sessionId]
  }

  function ensureSessionEntry(sessionId: string): SessionRecord {
    const existing = state.sessions[sessionId]
    if (existing) {
      return existing
    }

    const now = Date.now()
    const session: SessionRecord = {
      id: sessionId,
      createdAt: now,
      updatedAt: now,
      messageIds: [],
    }

    setState("sessions", sessionId, session)
    setState("sessionOrder", (order) => (order.includes(sessionId) ? order : [...order, sessionId]))
    return session
  }

  function addOrUpdateSession(input: SessionUpsertInput) {
    const session = ensureSessionEntry(input.id)
    const previousIds = [...session.messageIds]
    const nextMessageIds = Array.isArray(input.messageIds) ? input.messageIds : session.messageIds

    setState("sessions", input.id, {
      ...session,
      title: input.title ?? session.title,
      parentId: input.parentId ?? session.parentId ?? null,
      updatedAt: Date.now(),
      messageIds: nextMessageIds,
      revert: input.revert ?? session.revert ?? null,
    })

    if (Array.isArray(input.messageIds) && !areMessageIdListsEqual(previousIds, nextMessageIds)) {
      recomputeLastAssistantMessageId(input.id, nextMessageIds)
      bumpSessionRevision(input.id)
    }
  }

  function hydrateMessages(sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable<MessageInfo>) {
    if (!Array.isArray(inputs) || inputs.length === 0) return

    ensureSessionEntry(sessionId)

    const incomingIds = inputs.map((item) => item.id)

    const normalizedRecords: Record<string, MessageRecord> = {}
    const now = Date.now()

    inputs.forEach((input) => {
      const normalizedParts = normalizeParts(input.id, input.parts)
      const shouldBump = Boolean(input.bumpRevision || normalizedParts)
      const previous = state.messages[input.id]
      normalizedRecords[input.id] = {
        id: input.id,
        sessionId: input.sessionId,
        role: input.role,
        status: input.status,
        createdAt: input.createdAt ?? previous?.createdAt ?? now,
        updatedAt: input.updatedAt ?? now,
        isEphemeral: input.isEphemeral ?? previous?.isEphemeral ?? false,
        revision: previous ? previous.revision + (shouldBump ? 1 : 0) : 0,
        partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [],
        parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {},
      }
    })

    const infoList = infos ? Array.from(infos) : undefined
    const usageState = infoList ? rebuildUsageStateFromInfos(infoList) : state.usage[sessionId]

    const nextMessages: Record<string, MessageRecord> = { ...state.messages }
    const nextMessageInfoVersion: Record<string, number> = { ...state.messageInfoVersion }
    const nextPendingParts: Record<string, PendingPartEntry[]> = { ...state.pendingParts }
    const nextPermissionsByMessage: Record<string, Record<string, PermissionEntry>> = {
      ...state.permissions.byMessage,
    }

    Object.entries(normalizedRecords).forEach(([id, record]) => {
      nextMessages[id] = record
    })

    if (infoList) {
      for (const info of infoList) {
        const messageId = info.id as string
        messageInfoCache.set(messageId, info)
        const currentVersion = nextMessageInfoVersion[messageId] ?? 0
        nextMessageInfoVersion[messageId] = currentVersion + 1
      }
    }

    batch(() => {
      setState("messages", () => nextMessages)
      setState("messageInfoVersion", () => nextMessageInfoVersion)
      setState("pendingParts", () => nextPendingParts)
      setState("permissions", "byMessage", () => nextPermissionsByMessage)

      if (usageState) {
        setState("usage", sessionId, usageState)
      }

      setState("sessions", sessionId, (session) => ({
        ...session,
        messageIds: incomingIds,
        updatedAt: Date.now(),
      }))
      recomputeLastAssistantMessageId(sessionId, incomingIds)

      Object.values(normalizedRecords).forEach((record) => {
        maybeUpdateLatestTodoFromRecord(record)
      })

      bumpSessionRevision(sessionId)
    })
  }

  function insertMessageIntoSession(sessionId: string, messageId: string) {
    ensureSessionEntry(sessionId)
    setState("sessions", sessionId, "messageIds", (ids = []) => {
      if (ids.includes(messageId)) {
        return ids
      }
      return [...ids, messageId]
    })
  }

  function normalizeParts(messageId: string, parts: ClientPart[] | undefined) {
    if (!parts || parts.length === 0) {
      return null
    }
    const map: MessageRecord["parts"] = {}
    const ids: string[] = []

    parts.forEach((part, index) => {
      const id = ensurePartId(messageId, part, index)
      const cloned = clonePart(part)
      map[id] = {
        id,
        data: cloned,
        revision: 0,
      }
      ids.push(id)
    })

    return { map, ids }
  }

  function upsertMessage(input: MessageUpsertInput) {
    const normalizedParts = normalizeParts(input.id, input.parts)
    const shouldBump = Boolean(input.bumpRevision || normalizedParts)
    const now = Date.now()

    let nextRecord: MessageRecord | undefined

    setState("messages", input.id, (previous) => {
      const revision = previous ? previous.revision + (shouldBump ? 1 : 0) : 0
      const record: MessageRecord = {
        id: input.id,
        sessionId: input.sessionId,
        role: input.role,
        status: input.status,
        createdAt: input.createdAt ?? previous?.createdAt ?? now,
        updatedAt: input.updatedAt ?? now,
        isEphemeral: input.isEphemeral ?? previous?.isEphemeral ?? false,
        revision,
        partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [],
        parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {},
      }
      nextRecord = record
      return record
    })

    if (nextRecord) {
      maybeUpdateLatestTodoFromRecord(nextRecord)
    }

    insertMessageIntoSession(input.sessionId, input.id)
    flushPendingParts(input.id)
    recomputeLastAssistantMessageId(input.sessionId)
    bumpSessionRevision(input.sessionId)
  }

  function bufferPendingPart(entry: PendingPartEntry) {
    setState("pendingParts", entry.messageId, (list = []) => [...list, entry])
  }

  function clearPendingPartsForMessage(messageId: string) {
    setState("pendingParts", (prev) => {
      if (!prev[messageId]) {
        return prev
      }
      const next = { ...prev }
      delete next[messageId]
      return next
    })
  }

  function rebindPermissionForPart(messageId: string, partId: string, part: ClientPart) {
    if (!messageId || !partId || part.type !== "tool") {
      return
    }

    const toolCallId =
      (part as any).callID ??
      (part as any).callId ??
      (part as any).toolCallID ??
      (part as any).toolCallId ??
      undefined
    if (!toolCallId) {
      return
    }

    setState(
      "permissions",
      "byMessage",
      messageId,
      produce((draft) => {
        if (!draft) return
        const existing = draft[partId]
        for (const [key, entry] of Object.entries(draft)) {
          if (!entry || entry.partId) continue
          const permissionCallId =
            (entry.permission as any).tool?.callID ??
            (entry.permission as any).tool?.callId ??
            (entry.permission as any).callID ??
            (entry.permission as any).callId ??
            (entry.permission as any).toolCallID ??
            (entry.permission as any).toolCallId ??
            (entry.permission as any).metadata?.callID ??
            (entry.permission as any).metadata?.callId ??
            undefined
          if (permissionCallId !== toolCallId) continue
          if (!existing || existing.permission.id === entry.permission.id) {
            entry.partId = partId
            draft[partId] = entry
            delete draft[key]
          }
          break
        }
      }),
    )
  }

  function applyPartUpdate(input: PartUpdateInput) {
    const message = state.messages[input.messageId]
    if (!message) {
      bufferPendingPart({ messageId: input.messageId, part: input.part, receivedAt: Date.now() })
      return
    }

    const partId = ensurePartId(input.messageId, input.part, message.partIds.length)
    const cloned = clonePart(input.part)

    setState(
      "messages",
      input.messageId,
      produce((draft: MessageRecord) => {
        if (!draft.partIds.includes(partId)) {
          draft.partIds = [...draft.partIds, partId]
        }
        const existing = draft.parts[partId]
        const nextRevision = existing ? existing.revision + 1 : (cloned as any).version ?? 0
        draft.parts[partId] = {
          id: partId,
          data: cloned,
          revision: nextRevision,
        }
        draft.updatedAt = Date.now()
        if (input.bumpRevision ?? true) {
          draft.revision += 1
        }
      }),
    )

    rebindPermissionForPart(input.messageId, partId, cloned)

    if (isCompletedTodoPart(cloned)) {
      recordLatestTodoSnapshot(message.sessionId, {
        messageId: input.messageId,
        partId,
        timestamp: Date.now(),
      })
    }
  
    // Any part update can change the rendered height of the message
    // list, so we treat it as a session revision for scroll purposes.
    bumpSessionRevision(message.sessionId)
  }

  function applyPartDelta(input: {
    messageId: string
    partId: string
    field: string
    delta: string
    bumpRevision?: boolean
    bumpSessionRevision?: boolean
  }) {
    if (!input?.messageId || !input.partId || !input.field || typeof input.delta !== "string") {
      return
    }

    const message = state.messages[input.messageId]
    if (!message) {
      // Best-effort: drop deltas for unknown messages.
      return
    }

    let applied = false

    setState(
      "messages",
      input.messageId,
      produce((draft: MessageRecord) => {
        const entry = draft.parts[input.partId]
        if (!entry?.data) return
        const part = entry.data as any
        const currentValue = part?.[input.field]
        if (typeof currentValue === "string" || currentValue === undefined || currentValue === null) {
          part[input.field] = `${currentValue ?? ""}${input.delta}`
          applied = true
        }
        if (!applied) return
        entry.revision += 1
        draft.updatedAt = Date.now()
        if (input.bumpRevision ?? true) {
          draft.revision += 1
        }
      }),
    )

    if (applied && (input.bumpSessionRevision ?? true)) {
      bumpSessionRevision(message.sessionId)
    }
  }

  function removeMessage(messageId: string) {
    if (!messageId) return

    const record = state.messages[messageId]
    const sessionIds = new Set<string>()

    if (record?.sessionId) {
      sessionIds.add(record.sessionId)
    } else {
      Object.values(state.sessions).forEach((session) => {
        if (session.messageIds.includes(messageId)) {
          sessionIds.add(session.id)
        }
      })
    }

    clearRecordDisplayCacheForMessages(instanceId, [messageId])

    batch(() => {
      sessionIds.forEach((sessionId) => {
        setState("sessions", sessionId, "messageIds", (ids = []) => ids.filter((id) => id !== messageId))
      })

      setState("messages", (prev) => {
        if (!prev[messageId]) return prev
        const next = { ...prev }
        delete next[messageId]
        return next
      })

      setState("messageInfoVersion", (prev) => {
        if (!(messageId in prev)) return prev
        const next = { ...prev }
        delete next[messageId]
        return next
      })

      messageInfoCache.delete(messageId)

      setState("pendingParts", (prev) => {
        if (!prev[messageId]) return prev
        const next = { ...prev }
        delete next[messageId]
        return next
      })

      setState("permissions", "byMessage", (prev) => {
        if (!prev[messageId]) return prev
        const next = { ...prev }
        delete next[messageId]
        return next
      })

      sessionIds.forEach((sessionId) => {
        withUsageState(sessionId, (draft) => removeUsageEntry(draft, messageId))
        if (state.latestTodos[sessionId]?.messageId === messageId) {
          clearLatestTodoSnapshot(sessionId)
        }
        recomputeLastAssistantMessageId(sessionId)
        bumpSessionRevision(sessionId)
      })
    })
  }

  function removeMessagePart(messageId: string, partId: string) {
    if (!messageId || !partId) return
    const message = state.messages[messageId]
    if (!message) return

    clearRecordDisplayCacheForMessages(instanceId, [messageId])

    batch(() => {
      setState(
        "messages",
        messageId,
        produce((draft: MessageRecord) => {
          if (!draft.parts[partId] && !draft.partIds.includes(partId)) return
          draft.partIds = draft.partIds.filter((id) => id !== partId)
          delete draft.parts[partId]
          draft.updatedAt = Date.now()
          draft.revision += 1
        }),
      )

      setState("permissions", "byMessage", messageId, (prev) => {
        if (!prev || !prev[partId]) return prev
        const next = { ...prev }
        delete next[partId]
        return next
      })

      bumpSessionRevision(message.sessionId)
    })
  }


  function flushPendingParts(messageId: string) {
    const pending = state.pendingParts[messageId]
    if (!pending || pending.length === 0) {
      return
    }
    const now = Date.now()
    const validEntries = pending.filter((entry) => now - entry.receivedAt <= PENDING_PART_MAX_AGE_MS)
    if (validEntries.length === 0) {
      clearPendingPartsForMessage(messageId)
      return
    }
    validEntries.forEach((entry) => applyPartUpdate({ messageId, part: entry.part }))
    clearPendingPartsForMessage(messageId)
  }

  function replaceMessageId(options: ReplaceMessageIdOptions) {
    if (options.oldId === options.newId) return
    const existing = state.messages[options.oldId]
    if (!existing) return

    const cloned: MessageRecord = {
      ...existing,
      id: options.newId,
      isEphemeral: false,
      updatedAt: Date.now(),
      partIds: options.clearParts ? [] : existing.partIds,
      parts: options.clearParts ? {} : existing.parts,
    }

    setState("messages", options.newId, cloned)
    setState("messages", (prev) => {
      const next = { ...prev }
      delete next[options.oldId]
      return next
    })

    const affectedSessions = new Set<string>()

    Object.values(state.sessions).forEach((session) => {
      const index = session.messageIds.indexOf(options.oldId)
      if (index === -1) return
      setState("sessions", session.id, "messageIds", (ids) => {
        const next = [...ids]
        next[index] = options.newId
        return next
      })
      affectedSessions.add(session.id)
    })

    affectedSessions.forEach((sessionId) => {
      recomputeLastAssistantMessageId(sessionId)
      bumpSessionRevision(sessionId)
    })

    const infoEntry = messageInfoCache.get(options.oldId)
    if (infoEntry) {
      messageInfoCache.set(options.newId, infoEntry)
      messageInfoCache.delete(options.oldId)
      const version = state.messageInfoVersion[options.oldId] ?? 0
      setState("messageInfoVersion", options.newId, version)
      setState("messageInfoVersion", (prev) => {
        const next = { ...prev }
        delete next[options.oldId]
        return next
      })
    }

    const permissionMap = state.permissions.byMessage[options.oldId]
    if (permissionMap) {
      setState("permissions", "byMessage", options.newId, permissionMap)
      setState("permissions", (prev) => {
        const next = { ...prev }
        const nextByMessage = { ...next.byMessage }
        delete nextByMessage[options.oldId]
        next.byMessage = nextByMessage
        return next
      })
    }

    const questionMap = state.questions.byMessage[options.oldId]
    if (questionMap) {
      setState("questions", "byMessage", options.newId, questionMap)
      setState("questions", (prev) => {
        const next = { ...prev }
        const nextByMessage = { ...next.byMessage }
        delete nextByMessage[options.oldId]
        next.byMessage = nextByMessage
        return next
      })
    }

    const pending = state.pendingParts[options.oldId]
    if (pending) {
      setState("pendingParts", options.newId, pending)
    }
    clearPendingPartsForMessage(options.oldId)
    maybeUpdateLatestTodoFromRecord(cloned)
  }

  function setMessageInfo(messageId: string, info: MessageInfo) {
    if (!messageId) return
    messageInfoCache.set(messageId, info)
    const nextVersion = (state.messageInfoVersion[messageId] ?? 0) + 1
    setState("messageInfoVersion", messageId, nextVersion)
    updateUsageWithInfo(info)
  }

  function getMessageInfo(messageId: string) {
    void state.messageInfoVersion[messageId]
    return messageInfoCache.get(messageId)
  }

  function upsertPermission(entry: PermissionEntry) {
    const messageKey = entry.messageId ?? "__global__"
    const partKey = entry.partId ?? entry.permission?.id ?? "__global__"

    setState(
      "permissions",
      produce((draft) => {
        draft.byMessage[messageKey] = draft.byMessage[messageKey] ?? {}
        draft.byMessage[messageKey][partKey] = entry
        const existingIndex = draft.queue.findIndex((item) => item.permission.id === entry.permission.id)
        if (existingIndex === -1) {
          draft.queue.push(entry)
        } else {
          draft.queue[existingIndex] = entry
        }
        if (!draft.active || draft.active.permission.id === entry.permission.id) {
          draft.active = entry
        }
      }),
    )
  }

  function removePermission(permissionId: string) {
    setState(
      "permissions",
      produce((draft) => {
        draft.queue = draft.queue.filter((item) => item.permission.id !== permissionId)
        if (draft.active?.permission.id === permissionId) {
          draft.active = draft.queue[0] ?? null
        }
        Object.keys(draft.byMessage).forEach((messageKey) => {
          const partEntries = draft.byMessage[messageKey]
          Object.keys(partEntries).forEach((partKey) => {
            if (partEntries[partKey].permission.id === permissionId) {
              delete partEntries[partKey]
            }
          })
          if (Object.keys(partEntries).length === 0) {
            delete draft.byMessage[messageKey]
          }
        })
      }),
    )
  }

  function getPermissionState(messageId?: string, partId?: string) {
    const messageKey = messageId ?? "__global__"
    const partKey = partId ?? "__global__"
    const entry = state.permissions.byMessage[messageKey]?.[partKey]
    if (!entry) return null
    const active = state.permissions.active?.permission.id === entry.permission.id
    return { entry, active }
  }

  function upsertQuestion(entry: QuestionEntry) {
    const messageKey = entry.messageId ?? "__global__"
    const partKey = entry.partId ?? entry.request?.id ?? "__global__"

    setState(
      "questions",
      produce((draft) => {
        draft.byMessage[messageKey] = draft.byMessage[messageKey] ?? {}
        draft.byMessage[messageKey][partKey] = entry
        const existingIndex = draft.queue.findIndex((item) => item.request.id === entry.request.id)
        if (existingIndex === -1) {
          draft.queue.push(entry)
        } else {
          draft.queue[existingIndex] = entry
        }
        if (!draft.active || draft.active.request.id === entry.request.id) {
          draft.active = entry
        }
      }),
    )
  }

  function removeQuestion(requestId: string) {
    setState(
      "questions",
      produce((draft) => {
        draft.queue = draft.queue.filter((item) => item.request.id !== requestId)
        if (draft.active?.request.id === requestId) {
          draft.active = draft.queue[0] ?? null
        }
        Object.keys(draft.byMessage).forEach((messageKey) => {
          const partEntries = draft.byMessage[messageKey]
          Object.keys(partEntries).forEach((partKey) => {
            if (partEntries[partKey].request.id === requestId) {
              delete partEntries[partKey]
            }
          })
          if (Object.keys(partEntries).length === 0) {
            delete draft.byMessage[messageKey]
          }
        })
      }),
    )
  }

  function getQuestionState(messageId?: string, partId?: string) {
    const messageKey = messageId ?? "__global__"
    const partKey = partId ?? "__global__"
    const entry = state.questions.byMessage[messageKey]?.[partKey]
    if (!entry) return null
    const active = state.questions.active?.request.id === entry.request.id
    return { entry, active }
  }

  function pruneMessagesAfterRevert(sessionId: string, revertMessageId: string) {
    const session = state.sessions[sessionId]
    if (!session) return
    const stopIndex = session.messageIds.indexOf(revertMessageId)
    if (stopIndex === -1) return
    const removedIds = session.messageIds.slice(stopIndex)
    const keptIds = session.messageIds.slice(0, stopIndex)
    if (removedIds.length === 0) return

    setState("sessions", sessionId, "messageIds", keptIds)

    setState("messages", (prev) => {
      const next = { ...prev }
      removedIds.forEach((id) => delete next[id])
      return next
    })

    setState("messageInfoVersion", (prev) => {
      const next = { ...prev }
      removedIds.forEach((id) => delete next[id])
      return next
    })

    removedIds.forEach((id) => messageInfoCache.delete(id))

    setState("pendingParts", (prev) => {
      const next = { ...prev }
      removedIds.forEach((id) => {
        if (next[id]) delete next[id]
      })
      return next
    })

    setState("permissions", "byMessage", (prev) => {
      const next = { ...prev }
      removedIds.forEach((id) => {
        if (next[id]) delete next[id]
      })
      return next
    })

    setState("questions", "byMessage", (prev) => {
      const next = { ...prev }
      removedIds.forEach((id) => {
        if (next[id]) delete next[id]
      })
      return next
    })

    withUsageState(sessionId, (draft) => {
      removedIds.forEach((id) => removeUsageEntry(draft, id))
    })

    recomputeLastAssistantMessageId(sessionId, keptIds)
    bumpSessionRevision(sessionId)
  }

  function setSessionRevert(sessionId: string, revert?: SessionRecord["revert"] | null) {
    if (!sessionId) return
    ensureSessionEntry(sessionId)
    if (revert?.messageID) {
      pruneMessagesAfterRevert(sessionId, revert.messageID)
    }
    setState("sessions", sessionId, "revert", revert ?? null)
  }

  function getSessionRevert(sessionId: string) {
    return state.sessions[sessionId]?.revert ?? null
  }

  function makeScrollKey(sessionId: string, scope: string) {
    return `${sessionId}:${scope}`
  }

  function setScrollSnapshot(sessionId: string, scope: string, snapshot: Omit<ScrollSnapshot, "updatedAt">) {
    const key = makeScrollKey(sessionId, scope)
    setState("scrollState", key, { ...snapshot, updatedAt: Date.now() })
  }

  function getScrollSnapshot(sessionId: string, scope: string) {
    const key = makeScrollKey(sessionId, scope)
    return state.scrollState[key]
  }

   function clearSession(sessionId: string) {
     if (!sessionId) return

    const messageIds = Object.values(state.messages)
      .filter((record) => record.sessionId === sessionId)
      .map((record) => record.id)
 
    storeLog.info("Clearing session data", { instanceId, sessionId, messageCount: messageIds.length })
    clearRecordDisplayCacheForMessages(instanceId, messageIds)
 
    batch(() => {
      setState("messages", (prev) => {
        const next = { ...prev }
        messageIds.forEach((id) => delete next[id])
        return next
      })

      setState("messageInfoVersion", (prev) => {
        const next = { ...prev }
        messageIds.forEach((id) => delete next[id])
        return next
      })

      messageIds.forEach((id) => messageInfoCache.delete(id))

      setState("pendingParts", (prev) => {
        const next = { ...prev }
        messageIds.forEach((id) => {
          if (next[id]) delete next[id]
        })
        return next
      })

      setState("permissions", "byMessage", (prev) => {
        const next = { ...prev }
        messageIds.forEach((id) => {
          if (next[id]) delete next[id]
        })
        return next
      })

      setState("questions", "byMessage", (prev) => {
        const next = { ...prev }
        messageIds.forEach((id) => {
          if (next[id]) delete next[id]
        })
        return next
      })

      setState("usage", (prev) => {
        const next = { ...prev }
        delete next[sessionId]
        return next
      })

      setState("sessionRevisions", (prev) => {
        const next = { ...prev }
        delete next[sessionId]
        return next
      })

      setState("lastAssistantMessageIds", (prev) => {
        const next = { ...prev }
        delete next[sessionId]
        return next
      })

      setState("scrollState", (prev) => {
        const next = { ...prev }
        const prefix = `${sessionId}:`
        Object.keys(next).forEach((key) => {
          if (key.startsWith(prefix)) {
            delete next[key]
          }
        })
        return next
      })

      setState("sessions", sessionId, (current) => {
        if (!current) return current
        return { ...current, messageIds: [] }
      })

      setState("sessions", (prev) => {
        const next = { ...prev }
        delete next[sessionId]
        return next
      })

      setState("sessionOrder", (ids) => ids.filter((id) => id !== sessionId))
    })

    clearLatestTodoSnapshot(sessionId)
 
    hooks?.onSessionCleared?.(instanceId, sessionId)
  }

 
   function clearInstance() {
     messageInfoCache.clear()
      setState(reconcile(createInitialState(instanceId)))
    }
 
    return {

     instanceId,
     state,
     setState,
     addOrUpdateSession,
      hydrateMessages,
      upsertMessage,
      applyPartUpdate,
      applyPartDelta,
      removeMessage,
      removeMessagePart,
      bufferPendingPart,
      flushPendingParts,
     replaceMessageId,
     setMessageInfo,
     getMessageInfo,
      upsertPermission,
      removePermission,
      getPermissionState,
      upsertQuestion,
      removeQuestion,
      getQuestionState,

     setSessionRevert,
     getSessionRevert,
      rebuildUsage,
      getSessionUsage,
      setScrollSnapshot,
      getScrollSnapshot,
      getSessionRevision: getSessionRevisionValue,
      getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
      getLastAssistantMessageId: getLastAssistantMessageIdValue,
      getLastCompactionMessageIndex,
      getMessage: (messageId: string) => state.messages[messageId],
      getLatestTodoSnapshot: (sessionId: string) => state.latestTodos[sessionId],
      clearSession,
      clearInstance,
     }
   }