def String.words (s : String) :=
  s.split Char.isWhitespace |>.filter (¬·.isEmpty) |>.toArray
def String.lines (s : String) :=
  s.splitOn "\n" |>.filter (¬·.isEmpty) |>.toArray

-- TODO: remove trailing
def Array.unlines (xs : Array String) : String := Id.run do
  let mut out := ""
  for line in xs do
    out := out ++ line ++ "\n"
  return out

def ex := "The heat bloomed\tin December\nas the carnival  season\n        kicked into gear.\nNearly helpless with sun and glare, I avoided Rio's brilliant\nsidewalks\n  and glittering beaches,\npanting in dark   corners\nand waiting out the inverted southern summer."

#eval ex.words
#eval ex.lines
#eval IO.println ex.lines.unlines

def Array.replicate : Nat → α → Array α
| 0,   _ => #[]
| n+1, x => replicate n x |>.push x

theorem Array.replicate_size_eq_n : (Array.replicate n x).size = n := by
  induction n <;> simp [Array.replicate, *]

def String.replicate : Nat → Char → String
| 0,   _ => ""
| n+1, x => replicate n x |>.push x

def sep (xs space : Array String)
  : String := Id.run do
  let mut out := ""
  let mut i := 0
  while h : i < xs.size do
    out := out.append xs[i]
    if let some space := space[i]?
    then out := out.append space
    i := i + 1
  return out

def justify : Nat → Array String → String
| _,   #[word] => word
| width, words =>
    let wordsLength   := words.foldr (·.length + ·) 0
    let spacesNeeded  := width - wordsLength
    let numSpaces     := words.size - 1
    let smallSpaceLen := spacesNeeded / numSpaces
    let smallSpace    := String.replicate smallSpaceLen ' '
    let largeSpaces   := spacesNeeded % numSpaces
    let smallSpaces   := numSpaces - largeSpaces
    let largeSpace    := smallSpace.push ' '
    let spaces := Array.replicate largeSpaces largeSpace
               ++ Array.replicate smallSpaces smallSpace
    sep words spaces

#eval justify 1 "Hello, world, how is it going?".words
#eval justify 1 "Hello, world, how is it going?".words |>.length

def breakUp (width : Nat) (words : Array String) : Array (Array String) :=
  let (_, line, lines) := words.foldl
    (λ(left, line, lines) word ↦
      if word.length < left
      then (left - (word.length + 1), line.push word, lines)
      else (width - word.length, #[word], lines.push line))
    (width, #[], #[])
  lines.push line

def solve (width : Nat) :=
  Array.unlines ∘ Array.map (justify width) ∘ breakUp width ∘ String.words

#eval IO.println $ solve 40 ex

def usage := "Usage: justify WIDTH"
def fail : IO UInt32 := do IO.println usage; return 1

def main : List String → IO UInt32
| [width] =>
  if let some width := width.toNat?
  then do
    let mut text := ""
    let stdin ← IO.getStdin
    while true do
      let line ← stdin.getLine
      if line.isEmpty then break
      text := text ++ line
    IO.println $ solve width text
    return 0
  else fail
| _ => fail