#!/usr/bin/env python3

"""
Split an stl file.

The idea is to help post-processing stl files made with mapelia, so they
can be printed more easily. It does not modify the original file, but
creates two new files that end with "_N.stl" and "_S.stl"
(or "_head.stl" and "_tail.stl" if using the option --number).
"""

import sys
import os
import struct

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter as fmt

from maps import check_if_exists

red = lambda txt: '\x1b[31m%s\x1b[0m' % txt


def main():
    parser = ArgumentParser(description=__doc__, formatter_class=fmt)
    add = parser.add_argument  # shortcut
    add('file', help='stl file')
    add('-n', '--name', default='',
        help='output file (if empty, it is generated from the image file name)')
    add('--number', type=int, default=0,
        help='split by leaving a given number of triangles in the first file')
    add('--overwrite', action='store_true',
        help='do not check if the output files already exist')
    add('--ignore-check', action='store_true',
        help='go ahead even if the input file does not look like an stl')
    args = parser.parse_args()

    if not os.path.isfile(args.file):
        sys.exit('File %s does not exist.' % args.file)

    if not valid_stl(args.file):
        if args.ignore_check:
            print('Going ahead anyway as you used --ignore-check ...')
        else:
            sys.exit('Cancelling (use --ignore-check to force processing it).')

    if args.number == 0:
        def find_class(n, data):
            v1_z, _, _, v2_z, _, _, v3_z = struct.unpack('<7f', data[20:48])
            return 'N' if v1_z + v2_z + v3_z >= 0 else 'S'
    else:
        def find_class(n, data):
            return 'head' if n < args.number else 'tail'

    name = args.name or args.file.rsplit('.', 1)[0]
    for group, triangles in extract_triangles(args.file, find_class).items():
        output = '%s_%s.stl' % (name, group)
        if not args.overwrite:
            check_if_exists(output)
        write_stl(output, triangles)


def extract_triangles(fname, find_class):
    "Return dict with the triangles belonging to each class of stl file fname"
    print('Processing file %s ...' % fname)
    triangles = {}
    with open(fname, 'rb') as fin:
        fin.read(84)  # discard header + number of triangles
        size_triangle = 12 * 4 + 2  # 12 floats and 2 dummy bytes
        for i, data in enumerate(iter(lambda: fin.read(size_triangle), b'')):
            triangles.setdefault(find_class(i, data), []).append(data)
    return triangles


def write_stl(fname, triangles):
    "Write binary-encoded list of triangles as an stl file"
    print('Writing file %s ...' % fname)
    with open(fname, 'wb') as fout:
        fout.write(b'\0' * 80)  # header (empty)
        fout.write(struct.pack('<I', len(triangles)))  # number of triangles
        for t in triangles:
            fout.write(t)


def valid_stl(fname):
    "Return True if file looks like an stl, warn and return False otherwise"
    if not fname.endswith('.stl'):  # soft warning
        print(red('File %s does not end in ".stl". Is it really an stl file?' %
                  fname))

    ntriangles = struct.unpack('<I', open(fname, 'rb').read(84)[-4:])[0]
    size_content = os.path.getsize(fname) - 84  # 84 = header + number
    size_triangle = (1 * 3 + 3 * 3) * 4 + 2  # 4 bytes per float
    # 1 normal vector, 3 vertices (with 3 components each), 2 dummy bytes
    if size_content / size_triangle == ntriangles:
        return True
    else:
        print(red('File %s does not look like an stl file:\n'
                  '  It declares to have %d triangles but has capacity for %g.'
                  % (fname, ntriangles, size_content / size_triangle)))
        return False



if __name__ == '__main__':
    main()
