#!/usr/bin/env python3
"""Write the armour placement out as a Blockbench project.

    bb_export.py <armory.jar> [out.bbmodel]

Everything that positions a piece is drawn: the wearer, Armory's plates, this
mod's bezels, and - as the only unlocked cubes - the pieces whose transforms are
ours to change. Each of those carries its destination in its own name, in an
`[edit:...]` tag, so the file is the only thing `bb_import.py` needs.

The reference geometry is locked rather than merely present. A gem is aimed at a
bezel, and a bezel that can be nudged by accident is a bezel that stops being an
answer.
"""
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mcmodel as mc
import ARMOUR_EDITOR as ed
import blockbench as bb

REPO = ed.REPO
SOCKETS = ed.SOCKETS
SPACE = 'src/main/resources/packs/armory/data/cmmodular/miapi/modules/armor/space'

# hue per body part, so a group reads as a unit in the outliner
COLOURS = {'left_arm': 1, 'right_arm': 2, 'left_leg': 3, 'right_leg': 4,
           'body': 5, 'head': 6}


def box_cube(name, quads, matrix, pivot, colour, locked):
    """One drawn thing, already composed, as a Blockbench cube."""
    pts = np.vstack([q.pts for q in quads])
    pts = (np.column_stack([pts, np.ones(len(pts))]) @ np.asarray(matrix).T)[:, :3]
    pts = pts + np.asarray(pivot, float)
    centre, lo, hi, r = bb.oriented_box(pts, matrix)
    return bb.cube(name,
                   origin=bb.to_bb(centre),
                   lo=bb.to_bb(centre)[0] * 0 + bb.to_bb(centre) + lo,   # placeholder
                   hi=bb.to_bb(centre) + hi,
                   rotation=bb.matrix_to_euler_zyx(bb.rot_to_bb(r)),
                   colour=colour, locked=locked)


def main(jar, out):
    class A:
        pass
    args = A()
    args.scene = 'sockets'; args.jar = [jar]; args.variant = 'default'
    args.gem = 'medium'; args.chain = False; args.exact_merge = False
    args.layer1 = args.layer2 = 'miapi:item/armor/base/iron/layer_1'
    res = mc.Resources([jar, os.path.join(REPO, 'src/main/resources')])

    elements, outliner = [], []
    textures, tex_index = [], {}

    def texture(ref):
        """Embed a texture once, and hand back the index a face refers to it by."""
        if ref is None:
            return None
        if ref not in tex_index:
            try:
                img = res.image(ref)
            except (KeyError, OSError):
                tex_index[ref] = None
                return None
            tex_index[ref] = len(textures)
            textures.append(bb.texture_entry(ref.split('/')[-1], img, len(textures)))
        return tex_index[ref]

    def add(grp_children, name, quads, matrix, pivot, colour, locked):
        pts = np.vstack([q.pts for q in quads])
        pts = (np.column_stack([pts, np.ones(len(pts))]) @ np.asarray(matrix).T)[:, :3]
        pts = pts + np.asarray(pivot, float)
        centre, lo, hi, r = bb.oriented_box(pts, matrix)
        c = bb.to_bb(centre)
        for q in quads:
            texture(q.texture)
        faces = bb.faces_from_quads(quads, matrix, r, tex_index)
        cu = bb.cube(name, origin=c, lo=c + lo, hi=c + hi,
                     rotation=bb.matrix_to_euler_zyx(bb.rot_to_bb(r)),
                     colour=colour, locked=locked, faces=faces)
        elements.append(cu)
        grp_children.append(cu['uuid'])

    # ---------------------------------------------------------------- wearer
    body = []
    for part, (o, s, off, mirror) in mc.HUMANOID.items():
        if part == 'hat':
            continue
        q = mc.cube_quads(o, s, off, 0.0, mirror, (64, 64),
                          'minecraft:entity/player/wide/steve', 0)
        add(body, f'{part} (wearer)', q, np.eye(4), mc.PIVOTS[part], 8, True)
    outliner.append(bb.group('WEARER  (locked reference)', body))

    # ------------------------------------------------------------ the limbs
    chest = json.loads(res.read('data/tm_armory/miapi/modules/armor/chestplate.json'))['slots']
    pants = json.loads(res.read('data/tm_armory/miapi/modules/armor/pants.json'))['slots']
    slots = {'arm_left': chest['arm_left'], 'arm_right': chest['arm_right'],
             'leg_left': pants['leg_left'], 'leg_right': pants['leg_right']}
    gem_q = mc.model_quads(res.model('miapi:models/item/armor/gems/medium/'
                                     '[material.texture].json', 'default'), res)

    for part, slot, model in ed.PLATES:
        kids = []
        limb_tr = slots[slot]['transform']
        limb = mc.transform_matrix(limb_tr)
        piv = mc.PIVOTS[part]
        col = COLOURS[part]

        plate = mc.model_quads(res.model(f'miapi:models/item/armor/model/{model}/'
                                         '[material.texture].json', 'default'), res)
        add(kids, f'{slot} plate', plate, limb, piv, col, True)
        try:
            bez = mc.model_quads(res.model(f'cmmodular:models/item/armor/model/{slot}/'
                                           'socket/[material.texture].json', 'default'), res)
            add(kids, f'{slot} bezel', bez, limb, piv, col, True)
        except KeyError:
            pass

        path = os.path.join(REPO, SOCKETS, f'{slot}.json')
        with open(path) as fh:
            gem_tr = json.load(fh)['data']['replace']['slots']['gem']['transform']
        world = mc.merge(mc.transform_matrix(gem_tr), limb, True)
        add(kids, f'{slot} GEM  [edit:socket/{slot}#gem]', gem_q, world, piv, col, False)
        outliner.append(bb.group(part, kids, origin=bb.to_bb(piv)))

    # ------------------------------------------------------------ space gear
    helmet = json.loads(res.read('data/tm_armory/miapi/modules/armor/helmet.json'))['slots']
    # The `hat` slot is what puts a helmet on the head rather than the neck.
    for part, name, outer_tr, path_in_repo in (
            ('body', 'life_support_socket', chest['chest_back']['transform'],
             f'{SPACE}/life_support_socket.json'),
            ('head', 'helmet_socket', helmet['hat']['transform'],
             f'{SPACE}/helmet_socket.json')):
        kids = []
        piv = mc.PIVOTS[part]
        outer = mc.transform_matrix(outer_tr)
        with open(os.path.join(REPO, path_in_repo)) as fh:
            module = json.load(fh)
        for i, entry in enumerate(module['data']['replace']['model']):
            tr = entry['transform']
            if tr.get('origin') not in mc.BODY_PARTS:
                continue          # inventory model; MIAPI never draws it on a body
            q = mc.model_quads(res.model(entry['path'], 'default'), res)
            world = mc.merge(mc.transform_matrix(tr), outer, True)
            add(kids, f'{name}  [edit:space/{name}#model{i}]', q, world, piv,
                COLOURS[part], False)
        if kids:
            outliner.append(bb.group(f'{part} - {name}', kids, origin=bb.to_bb(piv)))

    doc = bb.project('cmmodular placement', outliner, elements, textures)
    with open(out, 'w') as fh:
        json.dump(doc, fh, indent=1)
    editable = [e['name'] for e in elements if not e['locked']]
    print(f'{out}\n  {len(elements)} cubes, {len(outliner)} groups, '
          f'{len(textures)} textures')
    print('  editable:')
    for n in editable:
        print(f'    {n}')


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else 'placement.bbmodel')