import type { Message, Part } from "@opencode-ai/sdk/v2/client"
import { mergeMessages } from "./optimistic"
import type { SessionMaterializationReason } from "./event-reducer"
import { sortMessagesChronologically } from "./message-ordering"

const STREAMING_PART_FIELDS = ["text", "output"] as const
const ACTIVE_TOOL_STATUSES = new Set(["pending", "running"])
const FINAL_TOOL_STATUSES = new Set(["completed", "error", "aborted", "failed", "timeout", "cancelled"])

export type MaterializedMessageRecord = {
  info: Message
  parts: Part[]
}

export type MaterializedState = {
  message: Record<string, Message[]>
  part: Record<string, Part[]>
}

export type MaterializeSessionSnapshotsOptions = {
  skipPartTypes?: ReadonlySet<string>
  mode?: "merge" | "prepend"
}

export type MaterializeSessionSnapshotsResult = {
  message: Record<string, Message[]>
  part: Record<string, Part[]>
  messages: Message[]
  messagesChanged: boolean
  partsChanged: boolean
}

export type SessionMaterializationStatus = {
  hasMessages: boolean
  renderable: boolean
  missingPartMessageIDs: string[]
}

export type SessionMaterializationRequest = {
  reason: SessionMaterializationReason
  messageID?: string
  partID?: string
}

export const getSessionMaterializationRequestKey = (
  runtimeKey: string,
  directory: string,
  sessionID: string,
): string => JSON.stringify([runtimeKey, directory, sessionID])

export function isSessionMaterializationStillNeeded(
  state: MaterializedState,
  sessionID: string,
  request: SessionMaterializationRequest,
): boolean {
  if (request.reason === "empty-assistant-message") {
    return !request.messageID || !Object.prototype.hasOwnProperty.call(state.part, request.messageID)
  }

  if (request.reason === "missing-owning-message") {
    if (!request.messageID) return true
    return !(state.message[sessionID] ?? []).some((message) => message.id === request.messageID)
  }

  if (request.reason === "orphan-delta" || request.reason === "missing-delta-part") {
    if (!request.messageID || !request.partID) return true
    return !(state.part[request.messageID] ?? []).some((part) => part.id === request.partID)
  }

  if (request.reason === "settled-running-tool") {
    return getStaleRunningToolMessageID(state, sessionID) === request.messageID
  }

  return true
}

export function getStaleRunningToolMessageID(
  state: MaterializedState,
  sessionID: string,
): string | undefined {
  const messages = state.message[sessionID] ?? []
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const message = messages[index]
    if (message.role === "user") return undefined
    if (message.role !== "assistant") continue
    const hasActiveTool = (state.part[message.id] ?? []).some((part) => {
      if (part.type !== "tool") return false
      const status = (part as { state?: { status?: unknown } }).state?.status
      return typeof status === "string" && ACTIVE_TOOL_STATUSES.has(status)
    })
    return hasActiveTool ? message.id : undefined
  }
  return undefined
}

function filterMaterializedParts(parts: Part[], skipPartTypes: ReadonlySet<string>): Part[] {
  return parts
    .filter((part) => !!part?.id && !skipPartTypes.has(part.type))
}

function haveEquivalentPartSnapshots(left: Part[] | undefined, right: Part[]): boolean {
  // `undefined` means "parts never fetched", which is NOT equivalent to a
  // fetched-empty snapshot — the empty array must be committed so
  // getSessionMaterializationStatus can tell the two apart.
  if (!left) return false
  if (left.length !== right.length) return false

  for (let index = 0; index < left.length; index += 1) {
    const leftPart = left[index]
    const rightPart = right[index]
    if (!leftPart || !rightPart) return false
    if (leftPart === rightPart) continue
    if (leftPart.id !== rightPart.id) return false
    if (JSON.stringify(leftPart) !== JSON.stringify(rightPart)) return false
  }

  return true
}

function getPartEndTime(part: Part): number | undefined {
  const stateEnd = (part as { state?: { time?: { end?: unknown } } }).state?.time?.end
  if (typeof stateEnd === "number") {
    return stateEnd
  }

  const timeEnd = (part as { time?: { end?: unknown } }).time?.end
  return typeof timeEnd === "number" ? timeEnd : undefined
}

function getStringField(part: Part, field: "text" | "output"): string | undefined {
  const value = (part as Record<string, unknown>)[field]
  return typeof value === "string" ? value : undefined
}

function getPartStateAttachments(part: Part): Array<unknown> | undefined {
  const state = (part as Record<string, unknown>).state as Record<string, unknown> | undefined
  if (!state) return undefined
  const attachments = state.attachments
  return Array.isArray(attachments) ? attachments : undefined
}

function hasLiveStreamingField(part: Part): boolean {
  if (getPartEndTime(part) !== undefined) return false
  return STREAMING_PART_FIELDS.some((field) => {
    const value = getStringField(part, field)
    return typeof value === "string" && value.length > 0
  })
}

function getPartStateTime(part: Part): { start?: number; end?: number } | undefined {
  const stateTime = (part as { state?: { time?: { start?: unknown; end?: unknown } } }).state?.time
  if (!stateTime || typeof stateTime !== "object") return undefined
  const start = typeof stateTime.start === "number" ? stateTime.start : undefined
  const end = typeof stateTime.end === "number" ? stateTime.end : undefined
  if (start === undefined && end === undefined) return undefined
  return { start, end }
}

function mergeMaterializedPart(existing: Part | undefined, next: Part): Part {
  if (!existing) return next

  if (existing.type === "tool" && next.type === "tool") {
    const existingStatus = (existing as { state?: { status?: unknown } }).state?.status
    const nextStatus = (next as { state?: { status?: unknown } }).state?.status
    if (
      typeof existingStatus === "string"
      && FINAL_TOOL_STATUSES.has(existingStatus)
      && typeof nextStatus === "string"
      && ACTIVE_TOOL_STATUSES.has(nextStatus)
    ) {
      return existing
    }
  }

  if (getPartEndTime(next) !== undefined) {
    const existingAttachments = getPartStateAttachments(existing)
    if (existingAttachments?.length && getPartStateAttachments(next) === undefined) {
      const nextRecord = { ...next }
      const nextState = { ...((next as Record<string, unknown>).state as Record<string, unknown> ?? {}), attachments: existingAttachments }
      ;(nextRecord as Record<string, unknown>).state = nextState
      return nextRecord
    }
    return next
  }

  let merged: Part = next
  for (const field of STREAMING_PART_FIELDS) {
    const existingValue = getStringField(existing, field)
    if (!existingValue) continue

    const nextValue = getStringField(next, field)
    if (typeof nextValue === "string" && nextValue.length >= existingValue.length) continue
    if (typeof nextValue === "string" && nextValue.length > 0 && !existingValue.startsWith(nextValue)) continue

    if (merged === next) merged = { ...next }
    const mergedRecord = merged as Record<string, unknown>
    mergedRecord[field] = existingValue
  }

  const existingAttachments = getPartStateAttachments(existing)
  if (existingAttachments?.length && getPartStateAttachments(next) === undefined) {
    if (merged === next) merged = { ...next }
    const mergedRecord = merged as Record<string, unknown>
    const nextState = (next as Record<string, unknown>).state as Record<string, unknown> | undefined
    const newState = { ...(nextState ?? {}), attachments: existingAttachments }
    mergedRecord.state = newState
  }

  const existingTime = getPartStateTime(existing)
  if (existingTime) {
    const nextTime = getPartStateTime(next)
    const preservedStart = nextTime?.start ?? existingTime.start
    const preservedEnd = nextTime?.end ?? existingTime.end
    if (preservedStart !== nextTime?.start || preservedEnd !== nextTime?.end) {
      if (merged === next) merged = { ...next }
      const mergedRecord = merged as Record<string, unknown>
      const currentState = (mergedRecord.state as Record<string, unknown> | undefined) ?? (next as Record<string, unknown>).state as Record<string, unknown> | undefined
      const newState = { ...(currentState ?? {}), time: { start: preservedStart, end: preservedEnd } }
      mergedRecord.state = newState
    }
  }

  return merged
}

function mergeMaterializedParts(
  existing: Part[] | undefined,
  nextParts: Part[],
  skipPartTypes: ReadonlySet<string>,
  preserveLiveStreamingParts: boolean,
): Part[] {
  if (!existing || existing.length === 0) return nextParts
  if (!preserveLiveStreamingParts) return nextParts

  const existingByID = new Map(existing.map((part) => [part.id, part]))
  let mergedParts = nextParts
  let changed = false

  for (let index = 0; index < nextParts.length; index += 1) {
    const nextPart = nextParts[index]
    const mergedPart = mergeMaterializedPart(existingByID.get(nextPart.id), nextPart)
    if (mergedPart === nextPart) continue
    if (!changed) mergedParts = [...nextParts]
    mergedParts[index] = mergedPart
    changed = true
  }

  const snapshotIDs = new Set(nextParts.map((part) => part.id))
  const missingLiveParts = existing.filter(
    (part) => !!part?.id && !snapshotIDs.has(part.id) && !skipPartTypes.has(part.type) && hasLiveStreamingField(part),
  )
  if (missingLiveParts.length === 0) return mergedParts

  return [...mergedParts, ...missingLiveParts]
}

export function materializeSessionSnapshots(
  state: MaterializedState,
  sessionID: string,
  records: MaterializedMessageRecord[],
  options: MaterializeSessionSnapshotsOptions = {},
): MaterializeSessionSnapshotsResult {
  const skipPartTypes = options.skipPartTypes ?? new Set<string>()
  const recordsByMessageID = new Map(
    records
      .filter((record) => !!record?.info?.id)
      .map((record) => [record.info.id, record] as const),
  )
  const nextMessages = sortMessagesChronologically([...recordsByMessageID.values()].map((record) => record.info))
  const snapshots = nextMessages.map((message) => recordsByMessageID.get(message.id)!)
  const existingMessages = state.message[sessionID]
  const currentMessages = existingMessages ?? []
  const incomingByID = new Map(nextMessages.map((message) => [message.id, message] as const))
  let reconciledCurrentMessages = currentMessages
  for (let index = 0; index < currentMessages.length; index += 1) {
    const existing = currentMessages[index]
    const incoming = incomingByID.get(existing.id)
    if (
      existing.role !== "assistant"
      || existing.error?.name !== "MessageAbortedError"
      || incoming?.role !== "assistant"
      || incoming.time.completed === undefined
    ) continue
    if (reconciledCurrentMessages === currentMessages) reconciledCurrentMessages = [...currentMessages]
    reconciledCurrentMessages[index] = incoming
  }
  const messages = mergeMessages(reconciledCurrentMessages, nextMessages)
  const messagesChanged = messages !== currentMessages || (existingMessages === undefined && snapshots.length === 0)

  let partsChanged = false
  let nextPartState = state.part
  const isPrepend = options.mode === "prepend"

  for (const record of snapshots) {
    const messageID = record.info.id
    if (isPrepend && nextPartState[messageID]) continue

    const isAssistant = record.info.role === "assistant"
    const existing = nextPartState[messageID]
    const nextParts = mergeMaterializedParts(
      existing,
      filterMaterializedParts(record.parts ?? [], skipPartTypes),
      skipPartTypes,
      isAssistant,
    )
    // For non-assistant messages an empty snapshot keeps the old "absent"
    // representation; only assistant messages need the explicit [] marker
    // (getSessionMaterializationStatus checks only assistant messages).
    const equivalent = existing
      ? haveEquivalentPartSnapshots(existing, nextParts)
      : nextParts.length === 0 && !isAssistant
    if (equivalent) continue

    if (nextPartState === state.part) nextPartState = { ...state.part }

    if (nextParts.length === 0 && !isAssistant) {
      delete nextPartState[messageID]
    } else {
      // Store fetched-empty as an explicit [] (not absence): an assistant
      // message the server returned with zero parts (e.g. aborted before any
      // output) is authoritatively empty and must count as renderable, or
      // the ensure-renderable effects retry syncSession forever.
      nextPartState[messageID] = nextParts
    }
    partsChanged = true
  }

  return {
    message: messagesChanged ? { ...state.message, [sessionID]: messages } : state.message,
    part: partsChanged ? nextPartState : state.part,
    messages,
    messagesChanged,
    partsChanged,
  }
}

export function getSessionMaterializationStatus(
  state: MaterializedState,
  sessionID: string,
): SessionMaterializationStatus {
  const messages = state.message[sessionID]
  if (!messages) {
    return { hasMessages: false, renderable: false, missingPartMessageIDs: [] }
  }

  const missingPartMessageIDs: string[] = []
  for (const message of messages) {
    if (message.role !== "assistant") continue
    // `undefined` = parts never fetched (not renderable yet). An explicit []
    // is a fetched-empty snapshot (e.g. aborted assistant turn) and counts
    // as renderable — otherwise sessions containing such a message can never
    // reach renderable state and ensure-renderable callers loop forever.
    const parts = state.part[message.id]
    if (!parts) {
      missingPartMessageIDs.push(message.id)
    }
  }

  return {
    hasMessages: true,
    renderable: missingPartMessageIDs.length === 0,
    missingPartMessageIDs,
  }
}