import std/[nre, sequtils, strutils, unicode]

# See
# - https://gist.github.com/ConnerWill/d4b6c776b509add763e17f9f113fd25b
# - https://en.wikipedia.org/wiki/ANSI_escape_code

const CSI = "\x1b["
const RESET_FORMATTING* = "\x1b[0m"
const ANSI_ESCAPE_CODE_ENDINGS*: seq[Rune] =
  toRunes(toSeq('A'..'Z') & toSeq('a'..'z'))

let FORMATTING_REGEX* = re("\x1b\\[([0-9;]*)([a-zA-Z])")

type
  TrueColor* = tuple[r, g, b: int]

  CursorType* = enum
    ctBlockBlink = 1, ctBlock, ctUnderlineBlink, ctUnderline, ctBarBlink, ctBar

  EraseMode* = enum
    emToEnd, emToStart, emAll

  TerminalColors* = enum
    cBlack = 0, cRed, cGreen, cYellow, cBlue, cMagenta, cCyan, cWhite,
    cDefault = 9

    cBrightBlack = 60, cBrightRed, cBrightGreen, cBrightYellow, cBrightBlue,
    cBrightMagenta, cBrightCyan, cBrightWhite

  TerminalStyle* = enum
    tsBold = 1, tsDim, tsItalic, tsUnderline, tsBlink, tsInverse, tsHidden,
    tsStrikethrough

proc stripFormatting*(text: string): string =
  text.replace(FORMATTING_REGEX, "")

proc stripAnsi*(text: string): string = stripFormatting(text)

func ansiAwareSubstring*(s: string, start, length: int): string =
  var outputRunes = newSeq[Rune]()
  var curAnsiSequences = newSeq[Rune]()
  var runes = s.toRunes
  var visibleRunesSeen = 0
  var visibleOutputLen = 0
  var i = 0

  while i < runes.len and visibleOutputLen < length:

    if runes.len > i + runeLen(CSI) and
       $runes[i..<i + len(CSI)] == CSI:
      # Recognize ANSI escape codes...

      var j = i + runeLen(CSI)
      while j < runes.len and runes[j] notin ANSI_ESCAPE_CODE_ENDINGS: j += 1

      # and remember it if we're before the start of the substring
      if visibleRunesSeen < start: curAnsiSequences.add(runes[i..j])

      # or add it to the output without increasing our substring length
      else: outputRunes.add(runes[i..j])

      # either way we want to pick up after it
      i = j

    else:
      if visibleRunesSeen >= start:
        outputRunes.add(runes[i])
        inc visibleOutputLen
      inc visibleRunesSeen

    inc i

    result = $(curAnsiSequences & outputRunes)


func ansiAwareSubstring*[T, U](s: string, r: HSlice[T,U]): string =
  let idxStart =
    when r.a is BackwardsIndex: s.runeLen - int(r.a)
    else: int(r.a)

  let idxEnd =
    when r.b is BackwardsIndex: s.runeLen - int(r.b)
    else: int(r.b)

  ansiAwareSubstring(s, idxStart, idxEnd - idxStart + 1)


func ansiEscSeq*(style: set[TerminalStyle]): string =
  CSI & toSeq(style).mapIt($int(it)).join(";") & "m"

func ansiEscSeq*(fg: int): string = CSI & "38;5;" & $fg & "m"
func ansiEscSeq*(bg: int): string = CSI & "48;5;" & $bg & "m"


func ansiEscSeq*(fg: TrueColor): string =
  "$#38;2;$#;$#;$#m" % [CSI, $fg.r, $fg.g, $fg.b]

func ansiEscSeq*(bg: TrueColor): string =
  "$#48;2;$#;$#;$#m" % [CSI, $bg.r, $bg.g, $bg.b]


func ansiEscSeq*(fg: TerminalColors): string = CSI & $(int(fg) + 30) & "m"
func ansiEscSeq*(bg: TerminalColors): string = CSI & $(int(bg) + 40) & "m"


func ansiEscSeq*(fg: TerminalColors, bg: TerminalColors, style: set[TerminalStyle]): string =
  result = CSI
  if style != {}: result &= toSeq(style).mapIt($int(it)).join(";") & ";"
  result &= $(int(fg) + 30) & ";" & $(int(bg) + 40) & "m"


func termFmt*(text: string, fg = cWhite, bg = cBlack, style: set[TerminalStyle] = {}): string =
  return ansiEscSeq(fg, bg, style) & text & RESET_FORMATTING


func termFmt*[C: int | TrueColor](
    text: string,
    fg, bg: C,
    style: set[TerminalStyle] = {}): string =
  return ansiEscSeq(style) & ansiEscSeq(fg = fg) & ansiEscSeq(bg = bg) &
    text & RESET_FORMATTING


func color*(text: string, fg = cDefault, bg = cDefault): string =
  return termFmt(text, fg, bg, {})

func invert*(text: string): string = return CSI & "7m" & text & RESET_FORMATTING
func striketrough*(text: string): string = return CSI & "9m" & text & RESET_FORMATTING


func eraseDisplay*(mode = emToEnd): string = return CSI & $int(mode) & "J"
func eraseLine*(mode = emToEnd): string = return CSI & $int(mode) & "K"


func blinkSlow*(text: string): string = return CSI & "5m" & text & RESET_FORMATTING
func blinkFast*(text: string): string = return CSI & "6m" & text & RESET_FORMATTING


func cursorUp*(n: int = 1): string = return CSI & $int(n) & "A"
func cursorDown*(n: int = 1): string = return CSI & $int(n) & "B"
func cursorForward*(n: int = 1): string = return CSI & $int(n) & "C"
func cursorBackward*(n: int = 1): string = return CSI & $int(n) & "D"
func cursorNextLine*(n: int = 1): string = return CSI & $int(n) & "E"
func cursorPrevLine*(n: int = 1): string = return CSI & $int(n) & "F"
func cursorHorizontalAbsolute*(col: int): string = return CSI & $int(col) & "G"


func cursorPosition*(row, col: int): string =
  return CSI & $int(row) & ";" & $int(col) & "H"


proc cursorType*(ct: CursorType): string = return CSI & $int(ct) & " q"
func saveCursorPosition*(): string = return CSI & "s"
func restoreCursorPosition*(): string = return CSI & "u"


func scrollUp*(n: int = 1): string = return CSI & $int(n) & "S"
func scrollDown*(n: int = 1): string = return CSI & $int(n) & "T"


proc write*(f: File, x, y: int, text: string) =
  f.write(cursorPosition(y, x) & text)


proc write*(x, y: int, text: string) =
  stdout.write(cursorPosition(y, x) & text)


proc setCursorType*(ct: CursorType) = stdout.write(cursorType(ct))
proc setCursorPosition*(x, y: int) = stdout.write(cursorPosition(y, x))
proc hideCursor*() = stdout.write(CSI & "?25l")
proc showCursor*() = stdout.write(CSI & "?25l")