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

def Array.unlines (xs : Array String) : String :=
  xs.foldl (· ++ · ++ "\n") ""

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

def sep (xs : Array String) (space : Fin (xs.size-1) → String)
  : String := Id.run do
  let mut out := ""
  let mut i := 0
  while h : i < xs.size do
    out := out ++ xs[i]
    if h : i < xs.size - 1 then
      out := out ++ space ⟨i, h⟩
    i := i + 1
  return out

def justify (width : Nat) (words : Array String) : String :=
  sep words space
  where
  space i := if i.val ≤ largeSpaces then largeSpace else smallSpace
  largeSpace    := smallSpace.push ' '
  smallSpace    := String.replicate smallSpaceLen ' '
  smallSpaceLen := spacesNeeded / numSpaces
  largeSpaces   := spacesNeeded % numSpaces
  numSpaces     := words.size - 1
  spacesNeeded  := width - wordsLength
  wordsLength   := words.foldr (·.length + ·) 0

def breakUp (width : Nat) (words : Array String) : Array (Array String) :=
  let (_, line, lines) := words.foldl
    (fun (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

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

def main : List String → IO UInt32
| [width] => do
  if let some width := width.toNat?
  then
    let stdin ← IO.getStdin
    let mut remainder : Array String := #[]
    repeat
      let line ← stdin.getLine
      if line.isEmpty then break
      let lines := breakUp width (remainder ++ line.words)
      if h : lines.size > 0 then
        remainder := lines[lines.size-1]
        for line in lines[0:lines.size-1] do
          IO.println <| justify width line
    IO.println <| sep remainder (fun _ => " ")
    return 0
  else fail
| _ => fail