import { useCallback, useState } from "react"
import { createClientMessageId } from "../../lib/chat-utils"
import { MAX_TRACE_ITEMS } from "../../lib/constants"

export type TraceTone = "info" | "error" | "success"

export interface ConversationTraceItem {
  id: string
  stage: string
  label: string
  detail?: string
  tone: TraceTone
  createdAt: string
}

export interface MessageTraceState {
  items: ConversationTraceItem[]
  expanded: boolean
}

function pushTraceItem(current: ConversationTraceItem[], item: ConversationTraceItem): ConversationTraceItem[] {
  const lastItem = current.at(-1)
  if (lastItem && lastItem.stage === item.stage && lastItem.label === item.label && lastItem.detail === item.detail && lastItem.tone === item.tone) {
    return current
  }

  const next = [...current, item]
  return next.slice(-MAX_TRACE_ITEMS)
}

export function useMessageTraces() {
  const [messageTraces, setMessageTraces] = useState<Record<string, MessageTraceState>>({})

  const renameMessageTrace = useCallback((currentMessageId: string, nextMessageId: string) => {
    setMessageTraces((current) => {
      if (!(currentMessageId in current)) {
        return current
      }
      const existing = current[currentMessageId]
      const { [currentMessageId]: _removed, ...rest } = current
      return {
        ...rest,
        [nextMessageId]: existing
      }
    })
  }, [])

  const initializeMessageTrace = useCallback((messageId: string, item?: ConversationTraceItem) => {
    setMessageTraces((current) => ({
      ...current,
      [messageId]: {
        items: item ? [item] : [],
        expanded: item?.tone === "error"
      }
    }))
  }, [])

  const appendTrace = useCallback((messageId: string, stage: string, label: string, detail?: string, tone: TraceTone = "info") => {
    setMessageTraces((current) => {
      const existing = current[messageId] ?? { items: [], expanded: false }
      return {
        ...current,
        [messageId]: {
          items: pushTraceItem(existing.items, {
            id: createClientMessageId("trace"),
            stage,
            label,
            detail,
            tone,
            createdAt: new Date().toISOString()
          }),
          expanded: tone === "error" ? true : existing.expanded
        }
      }
    })
  }, [])

  const setMessageTraceExpanded = useCallback((messageId: string, expanded: boolean) => {
    setMessageTraces((current) => {
      const existing = current[messageId]
      if (!existing || existing.expanded === expanded) {
        return current
      }

      return {
        ...current,
        [messageId]: {
          ...existing,
          expanded
        }
      }
    })
  }, [])

  const collapseMessageTrace = useCallback((messageId: string) => {
    setMessageTraceExpanded(messageId, false)
  }, [setMessageTraceExpanded])

  const removeMessageTrace = useCallback((messageId: string) => {
    setMessageTraces((current) => {
      if (!(messageId in current)) {
        return current
      }
      const { [messageId]: _removed, ...rest } = current
      return rest
    })
  }, [])

  return {
    messageTraces,
    renameMessageTrace,
    initializeMessageTrace,
    appendTrace,
    setMessageTraceExpanded,
    collapseMessageTrace,
    removeMessageTrace
  }
}
