#!/usr/bin/env python3

# Splits a file "X.lean" into "XSheet.lean" and "XSolution.lean".
# The script searches for sections of the form
#
#     --sol [Y]
#     [Z]
#     --los
#
# For each of these sections, "XSheet.lean" contains only [Y], whereas
# "XSolution.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"_DemoMaster", b"_Demo")
               .replace(b"_Exercise", b"_ExerciseSheet"))
      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])
