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

def List.unlines (xs : List String) : String
  := xs.intersperse "\n" |> String.join

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 sep : List a → List a → List a
| a1::a1s, a2::a2s => a1::a2::sep a1s a2s
| a1,       _      => a1

def justify : Nat → List String → String
| _,    [word] => word
| width, words =>
    let wordsLength   := List.sum (String.length <$> words)
    let spacesNeeded  := width - wordsLength
    let numSpaces     := words.length - 1
    let smallSpaceLen := spacesNeeded / numSpaces
    let smallSpace    := List.replicate smallSpaceLen ' '
    let largeSpaces   := spacesNeeded % numSpaces
    let smallSpaces   := numSpaces - largeSpaces
    let largeSpace    := ' '::smallSpace
    let spaces := List.replicate largeSpaces largeSpace
               ++ List.replicate smallSpaces smallSpace
               |>.map List.asString
    String.join $ 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 : List String) : List (List String) :=
  let (_, line, lines) := words.foldl
    (λ(left, line, lines) word ↦
      if word.length < left
      then (left - (word.length + 1), line++[word], lines)
      else (width - word.length, [word], lines++[line]))
    (width, [], [])
  lines++[line]

-- def breakUp (width : Nat) (words : List String)
--   : List (List String) :=
--   if width = 0 then [] else go width words [] []
--   where
--   go (left : Nat) (words line : List String) (lines : List $ List String) :=
--   match left, words, line, lines with
--   | _,    [],      l, ls => ls++[l]
--   | left, (w::ws), l, ls =>
--       if width < w.length
--       then go width ws [w] ls
--       else if left < w.length + 1
--       then go width (w::ws) [] (ls++[l])
--       else go (left - (w.length + 1)) ws (l++[w]) ls
--   termination_by (left, words)


def solve (width : Nat) :=
  List.unlines ∘ List.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