#!/usr/bin/env python3

# Splits a file "X.lean" into "X_sheet.lean" and "X_solution.lean".
# The script searches for sections of the form
#
#     --sol [Y]
#     [Z]
#     --los
#
# For each of these sections, "X_sheet.lean" contains only [Y], whereas
# "X_solution.lean" contains only [Z].

import os
import sys

def split_lean(name):
  (root, ext) = os.path.splitext(name)
  inp = open(name, "rb")
  sheet = open(root + "_sheet" + ext, "wb")
  solution = open(root + "_solution" + ext, "wb")

  in_solution = False
  line_number = 1
  for l in inp:
    if l.startswith(b"--sol ") or l.rstrip(b"\n\r") == b"--sol":
      if in_solution:
        print("Warning: missing --los on line {0} of {1}".format(line_number, name))
      else:
        in_solution = True
      if l[6:].startswith(b":="):
          sheet.seek(-1, os.SEEK_CUR)
          sheet.write(b" :=\n" + l[9:])
      else:
          sheet.write(l[6:])
    elif l.startswith(b"--los"):
      if not in_solution:
        print("Warning: missing --sol on line {0} of {1}".format(line_number, name))
      else:
        in_solution = False
    elif l.startswith(b"import "):
      newl = (l.replace(b"demo_master", b"demo")
               .replace(b"exercise", b"exercise_sheet")
               .replace(b"homework", b"homework_sheet"))
      sheet.write(newl)
      solution.write (newl)
    else:
      if not in_solution:
        sheet.write(l)
      solution.write(l)
    line_number += 1

if __name__ == "__main__":
    if len(sys.argv) <= 1 :
        print("Please provide the name of the file to process")
    else:
        split_lean(sys.argv[1])
