"""Step 2 of the controller-model pipeline: classify parts, move the meshes into WebXR grip
space and pack them into src/core/page/controllers/models.js.

    python3 bake.py <work folder from tessellate.js> <path to models.js>

Needs numpy.
"""
import base64
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from parts import PARTS, P, classify, face_stats, load_mesh  # noqa: E402

# CAD (mm, Valve's frame) -> WebXR grip space (m). Grip space: origin at the centre of the
# fist, -Z along the handle toward the thumb, +X out of the back of the right hand.
# The handle axis was fitted from the mesh (~53° from the faceplate); the origin is on that
# axis about level with the middle finger. Fine-tune in config.js (fit), not here.
GRIP_ORIGIN_MM = np.array([0.0, -50.0, -29.0])
CAD_TO_GRIP = np.array([[1, 0, 0],          # grip X
                        [0, -0.8, 0.6],     # grip Y
                        [0, -0.6, -0.8]])   # grip Z
STICK_PIVOT_MM = {'R': [-13.1, -10.2, -3.0], 'L': [13.1, -10.2, -3.0]}
SCALE = 20000.0                             # Int16 units per metre (0.05 mm steps, ±1.6 m range)

HEADER = '''// Steam Frame controller models for "StreamXR".
// GENERATED by tools/controller-models/bake.py - do not edit by hand.
//
// Adapted from Valve's Steam Frame controller CAD, (c) 2026 Valve Corporation.
// Source: https://gitlab.steamos.cloud/SteamHardware/SteamFrame (folder documents-controller)
// Modified: the STEP files were tessellated into triangle meshes, split into separate parts
// (buttons, trigger, grip, etc.), transformed into WebXR grip space, and packed into this file.
// License: CC BY-NC-SA 4.0 (non-commercial) https://creativecommons.org/licenses/by-nc-sa/4.0/
// Provided as-is, without warranty (see the license).
// This file is NOT covered by the repo's MIT license (see LICENSE and README.md).
// Not affiliated with, endorsed by, or sponsored by Valve Corporation.
//
// Layout per hand: Int16 pos[verts*3] (metres*scale, WebXR grip space) | Int8 normal[verts*3] (/127)
//                  | Uint8 part[verts] | pad | Uint16 index[tris*3]
'''


def to_grip(p_mm):
    # Point(s) in CAD millimetres -> grip space in metres: shift to the grip origin, rotate
    # into grip axes, convert mm to m.
    return (np.asarray(p_mm) - GRIP_ORIGIN_MM) @ CAD_TO_GRIP.T / 1000.0


def anchors(fine_path, face_parts, side):
    """Per part: centre, average outward normal and highest point (grip space). Used to
    place icons and to know where the stick pivots."""
    v, _, tris, face = load_mesh(fine_path)
    t = v[tris]                                          # the three corners of every triangle
    c = t.mean(1)                                        # triangle centres
    cr = np.cross(t[:, 1] - t[:, 0], t[:, 2] - t[:, 0])  # normal * 2*area, per triangle
    area = np.linalg.norm(cr, axis=1)
    tri_part = face_parts[face]                          # which part each triangle belongs to
    out = {}
    for k, name in enumerate(PARTS):
        sel = tri_part == k
        if not sel.any():
            continue
        # Area-weighted centre, summed normal (so large faces count more), and the point
        # furthest out along that normal (the part's "top").
        centre = (c[sel] * area[sel, None]).sum(0) / area[sel].sum()
        normal = cr[sel].sum(0)
        normal /= np.linalg.norm(normal)
        top = c[sel][np.argmax(c[sel] @ normal)]
        out[name] = {'c': to_grip(centre).round(5).tolist(), 'n': (CAD_TO_GRIP @ normal).round(4).tolist(),
                     'top': to_grip(top).round(5).tolist()}
    # The thumbstick tilts about a fixed point (measured on the CAD); its "up" direction is
    # the faceplate normal (CAD +Z), turned into grip space.
    out['stickPivot'] = {'c': to_grip(STICK_PIVOT_MM[side]).round(5).tolist(),
                         'n': (CAD_TO_GRIP @ np.array([0, 0, 1.0])).round(4).tolist()}
    return out


def bake_hand(work, side):
    # Classify every CAD face on the fine mesh, then copy those labels onto the coarse mesh
    # (the one that ships); both come from the same CAD faces, in the same order.
    fine = work / f'fine_{side}.json'
    face_parts = classify(fine, side, work / 'fine_R.json')

    v, n, tris, face = load_mesh(work / f'coarse_{side}.json')
    if face.max() + 1 != len(face_parts):
        raise SystemExit(f'{side}: coarse and fine meshes have different CAD faces')
    vpart = np.zeros(len(v), np.uint8)                   # OCCT gives each CAD face its own vertices
    vpart[tris.ravel()] = np.repeat(face_parts[face], 3)

    # Append the battery door as extra geometry (its triangle indices shift by the body's
    # vertex count) and label all of it as the 'door' part.
    dv, dn, dtris, _ = load_mesh(work / f'door_{side}.json')
    V = np.vstack([v, dv])
    N = np.vstack([n, dn])
    T = np.vstack([tris, dtris + len(v)])
    parts = np.r_[vpart, np.full(len(dv), P['door'], np.uint8)]

    # Move positions and normals into grip space; refuse meshes the 16-bit encoding can't hold.
    G = to_grip(V)
    NG = N @ CAD_TO_GRIP.T
    NG /= np.linalg.norm(NG, axis=1, keepdims=True) + 1e-9
    if np.abs(G).max() * SCALE >= 32767 or len(V) >= 65536:
        raise SystemExit(f'{side}: mesh too big for the Int16/Uint16 encoding')

    # Pack as documented in HEADER: positions (Int16), normals (Int8), part per vertex
    # (Uint8), a pad byte if needed so the Uint16 indices start on an even offset, indices.
    blob = np.round(G * SCALE).astype('<i2').tobytes() + np.round(NG * 127).astype('i1').tobytes() + parts.tobytes()
    pad = (-len(blob)) % 2
    blob += b'\0' * pad + T.astype('<u2').tobytes()
    print(f'{side}: {len(V)} vertices, {len(T)} triangles, {len(blob) // 1024} KB')
    return {'verts': int(len(V)), 'tris': int(len(T)), 'scale': SCALE, 'pad': pad,
            'anchors': anchors(fine, face_parts, side), 'data': base64.b64encode(blob).decode()}


def main():
    if len(sys.argv) != 3:
        raise SystemExit(__doc__)
    work, out = Path(sys.argv[1]), Path(sys.argv[2])
    data = {'parts': PARTS, 'R': bake_hand(work, 'R'), 'L': bake_hand(work, 'L')}
    out.write_text(HEADER + '(globalThis.__streamxr = globalThis.__streamxr || {}).controllerModels = '
                   + json.dumps(data) + ';\n', encoding='utf-8')
    print(f'wrote {out}')


if __name__ == '__main__':
    main()