#!/usr/bin/env python

"""
Convert additional staves in musicxml scores into ossias. This is useful for making scores
searchable by search.py.

In particular, this script takes a 2-staff file and adds the following code to it:
<staff-details number="1">
  <staff-type>ossia</staff-type>
</staff-details>

This marks staff #1 as an ossia, so that it is recognized as a variant by search.py, so search.py
can find substrings in and around the variant.

Installation:
* Clone the repo
* Make a virtualenv with python3 and pip install lxml
"""


from argparse import ArgumentParser
from io import BytesIO
from zipfile import ZipFile

# There's no way to preserve xml and doctype declarations with xml.etree
from lxml import etree as ET


def parse_args():
    parser = ArgumentParser(description="""
        Convert a 2-staff musicxml score into a score with a single main staff and an ossia staff,
        as annotated by the <staff-details> element. The convention followed by this script is
        that the first staff is the ossia staff and the second is the main staff.
    """)

    parser.add_argument('file', help="File to modify")
    parser.add_argument('-t', '--test', action='store_true',
        help="Show the modified musicxml instead of saving it")

    return parser.parse_args()


def add_ossia_marker(tree):
    attributes = tree.getroot().find('./part/measure/attributes')

    num_staves = int(attributes.find('staves').text)
    if num_staves != 2:
        raise ValueError(f"Expected only 2 staves, got {num_staves}")

    last_clef_idx = [i for i, el in enumerate(attributes) if el.tag == 'clef'][-1]
    attributes.insert(last_clef_idx + 1, ET.fromstring(
        '<staff-details number="1"><staff-type>ossia</staff-type></staff-details>'
    ))


def tree2str(tree):
    # This appears to be the only way to preserve the xml and doctype declarations
    buffer = BytesIO()
    tree.write(buffer, xml_declaration=True, encoding=tree.docinfo.encoding)
    return buffer.getvalue().decode('utf8')


def main():
    args = parse_args()

    with ZipFile(args.file) as zipf:
        maybe_fnames = [f for f in zipf.namelist() if not f.startswith('META-INF/')]
        assert len(maybe_fnames) == 1, f"More than one relevant file found in {args.file}! Malformed mxl file?"
        fname = maybe_fnames[0]

        with zipf.open(fname) as f:
            tree = ET.parse(f)

    add_ossia_marker(tree)

    if args.test:
        print(tree2str(tree))
    else:
        # NB: this drops the META-INF member of the zip archive. Nobody cares.
        with ZipFile(args.file, 'w') as zipf:
            zipf.writestr(fname, tree2str(tree))



if __name__ == '__main__':
    main()