1  import {
  2    type AnsiCode,
  3    ansiCodesToString,
  4    reduceAnsiCodes,
  5    type Token,
  6    tokenize,
  7    undoAnsiCodes,
  8  } from '@alcalzone/ansi-tokenize'
  9  import type { Theme } from './theme.js'
 10  
 11  export type TextHighlight = {
 12    start: number
 13    end: number
 14    color: keyof Theme | undefined
 15    dimColor?: boolean
 16    inverse?: boolean
 17    shimmerColor?: keyof Theme
 18    priority: number
 19  }
 20  
 21  export type TextSegment = {
 22    text: string
 23    start: number
 24    highlight?: TextHighlight
 25  }
 26  
 27  export function segmentTextByHighlights(
 28    text: string,
 29    highlights: TextHighlight[],
 30  ): TextSegment[] {
 31    if (highlights.length === 0) {
 32      return [{ text, start: 0 }]
 33    }
 34  
 35    const sortedHighlights = [...highlights].sort((a, b) => {
 36      if (a.start !== b.start) return a.start - b.start
 37      return b.priority - a.priority
 38    })
 39  
 40    const resolvedHighlights: TextHighlight[] = []
 41    const usedRanges: Array<{ start: number; end: number }> = []
 42  
 43    for (const highlight of sortedHighlights) {
 44      if (highlight.start === highlight.end) continue
 45  
 46      const overlaps = usedRanges.some(
 47        range =>
 48          (highlight.start >= range.start && highlight.start < range.end) ||
 49          (highlight.end > range.start && highlight.end <= range.end) ||
 50          (highlight.start <= range.start && highlight.end >= range.end),
 51      )
 52  
 53      if (!overlaps) {
 54        resolvedHighlights.push(highlight)
 55        usedRanges.push({ start: highlight.start, end: highlight.end })
 56      }
 57    }
 58  
 59    return new HighlightSegmenter(text).segment(resolvedHighlights)
 60  }
 61  
 62  class HighlightSegmenter {
 63    private readonly tokens: Token[]
 64    // Two position systems: "visible" (what the user sees, excluding ANSI codes)
 65    // and "string" (raw positions including ANSI codes for substring extraction)
 66    private visiblePos = 0
 67    private stringPos = 0
 68    private tokenIdx = 0
 69    private charIdx = 0 // offset within current text token (for partial consumption)
 70    private codes: AnsiCode[] = []
 71  
 72    constructor(private readonly text: string) {
 73      this.tokens = tokenize(text)
 74    }
 75  
 76    segment(highlights: TextHighlight[]): TextSegment[] {
 77      const segments: TextSegment[] = []
 78  
 79      for (const highlight of highlights) {
 80        const before = this.segmentTo(highlight.start)
 81        if (before) segments.push(before)
 82  
 83        const highlighted = this.segmentTo(highlight.end)
 84        if (highlighted) {
 85          highlighted.highlight = highlight
 86          segments.push(highlighted)
 87        }
 88      }
 89  
 90      const after = this.segmentTo(Infinity)
 91      if (after) segments.push(after)
 92  
 93      return segments
 94    }
 95  
 96    private segmentTo(targetVisiblePos: number): TextSegment | null {
 97      if (
 98        this.tokenIdx >= this.tokens.length ||
 99        targetVisiblePos <= this.visiblePos
100      ) {
101        return null
102      }
103  
104      const visibleStart = this.visiblePos
105  
106      // Consume leading ANSI codes before first visible char
107      while (this.tokenIdx < this.tokens.length) {
108        const token = this.tokens[this.tokenIdx]!
109        if (token.type !== 'ansi') break
110        this.codes.push(token)
111        this.stringPos += token.code.length
112        this.tokenIdx++
113      }
114  
115      const stringStart = this.stringPos
116      const codesStart = [...this.codes]
117  
118      // Advance through tokens until we reach target
119      while (
120        this.visiblePos < targetVisiblePos &&
121        this.tokenIdx < this.tokens.length
122      ) {
123        const token = this.tokens[this.tokenIdx]!
124  
125        if (token.type === 'ansi') {
126          this.codes.push(token)
127          this.stringPos += token.code.length
128          this.tokenIdx++
129        } else {
130          const charsNeeded = targetVisiblePos - this.visiblePos
131          const charsAvailable = token.value.length - this.charIdx
132          const charsToTake = Math.min(charsNeeded, charsAvailable)
133  
134          this.stringPos += charsToTake
135          this.visiblePos += charsToTake
136          this.charIdx += charsToTake
137  
138          if (this.charIdx >= token.value.length) {
139            this.tokenIdx++
140            this.charIdx = 0
141          }
142        }
143      }
144  
145      // Empty segment (can occur when only trailing ANSI codes remain)
146      if (this.stringPos === stringStart) {
147        return null
148      }
149  
150      const prefixCodes = reduceCodes(codesStart)
151      const suffixCodes = reduceCodes(this.codes)
152      this.codes = suffixCodes
153  
154      const prefix = ansiCodesToString(prefixCodes)
155      const suffix = ansiCodesToString(undoAnsiCodes(suffixCodes))
156  
157      return {
158        text: prefix + this.text.substring(stringStart, this.stringPos) + suffix,
159        start: visibleStart,
160      }
161    }
162  }
163  
164  function reduceCodes(codes: AnsiCode[]): AnsiCode[] {
165    return reduceAnsiCodes(codes).filter(c => c.code !== c.endCode)
166  }