"""
A minimal glTF 2.0 binary writer: one mesh, one PBR material, embedded
textures. Written by hand rather than through trimesh because the remesher
needs to store TANGENT: the normal map is baked in exactly these tangents, and
a renderer that derives its own would disagree with the bake at every
triangle, which shows as faceting on a low-poly model.
"""

import io
import json
import struct

import numpy as np
from PIL import Image

FLOAT, UINT32 = 5126, 5125
ARRAY_BUFFER, ELEMENT_ARRAY_BUFFER = 34962, 34963


def _encode(image: Image.Image, photo: bool) -> tuple[bytes, str]:
    """Colour without transparency goes as JPEG (a tenth of the size, and it is a
    photograph-like texture); data maps stay PNG, where JPEG artefacts would
    read as bumps and shine."""
    out = io.BytesIO()
    if photo and (image.mode != "RGBA" or image.getchannel("A").getextrema()[0] == 255):
        image.convert("RGB").save(out, format="JPEG", quality=92)
        return out.getvalue(), "image/jpeg"
    image.save(out, format="PNG", compress_level=6)
    return out.getvalue(), "image/png"


def write_glb(path, positions, normals, tangents, uvs, indices, textures, metallic=1.0, roughness=1.0, generator="3dvalley.com remesh"):
    """
    positions, normals: (n, 3); tangents: (n, 4) with handedness in w;
    uvs: (n, 2) with v up (flipped to glTF's v down here); indices: (m, 3).
    textures: {"color": Image, "metallicRoughness": Image | None, "normal": Image}.
    Without a normal map (textures["normal"] None), tangents may be None and
    are then left out of the file.
    """
    positions = np.ascontiguousarray(positions, dtype=np.float32)
    normals = np.ascontiguousarray(normals, dtype=np.float32)
    tangents = np.ascontiguousarray(tangents, dtype=np.float32) if tangents is not None else None
    uvs = np.ascontiguousarray(np.stack([uvs[:, 0], 1.0 - uvs[:, 1]], axis=1), dtype=np.float32)
    indices = np.ascontiguousarray(indices, dtype=np.uint32).ravel()

    blob = bytearray()
    views, accessors = [], []

    def add_view(data: bytes, target=None):
        while len(blob) % 4:
            blob.append(0)
        view = {"buffer": 0, "byteOffset": len(blob), "byteLength": len(data)}
        if target:
            view["target"] = target
        blob.extend(data)
        views.append(view)
        return len(views) - 1

    def add_accessor(array, kind, component, target, with_bounds=False):
        view = add_view(array.tobytes(), target)
        accessor = {"bufferView": view, "componentType": component, "count": int(array.shape[0]), "type": kind}
        if with_bounds:
            accessor["min"] = array.min(axis=0).tolist()
            accessor["max"] = array.max(axis=0).tolist()
        accessors.append(accessor)
        return len(accessors) - 1

    attributes = {
        "POSITION": add_accessor(positions, "VEC3", FLOAT, ARRAY_BUFFER, with_bounds=True),
        "NORMAL": add_accessor(normals, "VEC3", FLOAT, ARRAY_BUFFER),
    }
    if tangents is not None:
        attributes["TANGENT"] = add_accessor(tangents, "VEC4", FLOAT, ARRAY_BUFFER)
    attributes["TEXCOORD_0"] = add_accessor(uvs, "VEC2", FLOAT, ARRAY_BUFFER)
    index_view = add_view(indices.tobytes(), ELEMENT_ARRAY_BUFFER)
    accessors.append({"bufferView": index_view, "componentType": UINT32, "count": int(indices.size), "type": "SCALAR"})
    index_accessor = len(accessors) - 1

    images, gltf_textures = [], []

    def add_texture(image, photo=False):
        data, mime = _encode(image, photo)
        view = add_view(data)
        images.append({"bufferView": view, "mimeType": mime})
        gltf_textures.append({"source": len(images) - 1, "sampler": 0})
        return len(gltf_textures) - 1

    pbr = {"baseColorTexture": {"index": add_texture(textures["color"], photo=True)}, "metallicFactor": metallic, "roughnessFactor": roughness}
    if textures.get("metallicRoughness") is not None:
        pbr["metallicRoughnessTexture"] = {"index": add_texture(textures["metallicRoughness"])}
    else:
        pbr["metallicFactor"] = 0.0
    material = {"pbrMetallicRoughness": pbr}
    if textures.get("normal") is not None:
        material["normalTexture"] = {"index": add_texture(textures["normal"])}

    while len(blob) % 4:
        blob.append(0)
    document = {
        "asset": {"version": "2.0", "generator": generator},
        "scene": 0,
        "scenes": [{"nodes": [0]}],
        "nodes": [{"mesh": 0, "name": "remeshed"}],
        "meshes": [{"primitives": [{"attributes": attributes, "indices": index_accessor, "material": 0}]}],
        "materials": [material],
        "textures": gltf_textures,
        "images": images,
        "samplers": [{"magFilter": 9729, "minFilter": 9987, "wrapS": 10497, "wrapT": 10497}],
        "buffers": [{"byteLength": len(blob)}],
        "bufferViews": views,
        "accessors": accessors,
    }
    text = json.dumps(document, separators=(",", ":")).encode()
    text += b" " * (-len(text) % 4)
    with open(path, "wb") as f:
        f.write(struct.pack("<III", 0x46546C67, 2, 12 + 8 + len(text) + 8 + len(blob)))
        f.write(struct.pack("<II", len(text), 0x4E4F534A) + text)
        f.write(struct.pack("<II", len(blob), 0x004E4942) + bytes(blob))


def write_obj(path, positions, polygons):
    """The new topology as OBJ, quads kept as quads for editors that care."""
    with open(path, "w") as f:
        f.write("# remeshed with 3dvalley.com\n")
        for x, y, z in positions:
            f.write(f"v {x:.6f} {y:.6f} {z:.6f}\n")
        for poly in polygons:
            unique = [int(i) + 1 for k, i in enumerate(poly) if i not in poly[:k]]
            f.write("f " + " ".join(map(str, unique)) + "\n")