# glTF 2.0 writer. Produces both .gltf+.bin pairs and .glb single files.
# Supports: positions, normals, tangents, UV0, vertex colors, indices,
# full PBR base color + metallic-roughness + normal + emissive + occlusion,
# hierarchy via parent indexing, and texture embedding for GLB.

import json
import os
import struct

from .. import settings
from . import texture_utils


def _encode_uri(uri):
    if not uri:
        return uri
    return texture_utils.encode_uri(uri)


class GltfWriter:
    def __init__(self):
        self._buffer = bytearray()
        self._buffer_views = []
        self._accessors = []
        self._meshes = []
        self._nodes = []
        self._materials = []
        self._images = []
        self._textures = []
        self._samplers = []
        self._lights = []
        self._cameras = []
        self._extensions_used = set()
        self._scene_root_indices = []

        # Caches
        self._texture_cache_by_uri = {}     # external mode
        self._texture_cache_by_path = {}    # embed mode
        self._material_cache_by_name = {}

        # Mode: "external" writes URIs to disk-referenced files; "embed" packs
        # bytes into the buffer for GLB.
        self._embed_textures = False

    def set_embed_textures(self, embed):
        self._embed_textures = bool(embed)

    # ---------- public assembly ----------

    def add_light(self, light_data):
        # light_data is a glTF KHR_lights_punctual light dict. Returns index.
        index = len(self._lights)
        self._lights.append(light_data)
        self._extensions_used.add("KHR_lights_punctual")
        return index

    def add_camera(self, camera_data):
        # camera_data is a core glTF perspective or orthographic camera dict.
        index = len(self._cameras)
        self._cameras.append(camera_data)
        return index

    def add_node(self, name, translation=None, rotation=None, scale=None,
                 mesh_index=None, parent_index=None, matrix=None,
                 light_index=None, camera_index=None):
        # Returns the new node's index. parent_index can be None (root).
        node = {"name": name}
        if mesh_index is not None:
            node["mesh"] = mesh_index
        if light_index is not None:
            node.setdefault("extensions", {})["KHR_lights_punctual"] = {"light": light_index}
        if camera_index is not None:
            node["camera"] = camera_index
        if matrix:
            node["matrix"] = [float(v) for v in matrix]
        else:
            if translation and any(abs(v) > 1e-9 for v in translation):
                node["translation"] = [float(v) for v in translation]
            if rotation and not (
                    abs(rotation[0]) < 1e-9
                    and abs(rotation[1]) < 1e-9
                    and abs(rotation[2]) < 1e-9
                    and abs(rotation[3] - 1.0) < 1e-9):
                node["rotation"] = [float(v) for v in rotation]
            if scale and any(abs(v - 1.0) > 1e-9 for v in scale):
                node["scale"] = [float(v) for v in scale]
        index = len(self._nodes)
        self._nodes.append(node)
        if parent_index is None:
            self._scene_root_indices.append(index)
        else:
            parent = self._nodes[parent_index]
            parent.setdefault("children", []).append(index)
        return index

    def add_mesh(self, mesh_name, primitives_payload):
        # primitives_payload is a list of dicts each containing primitive data
        # and a 'material' dict (see _add_material) or None.
        gltf_primitives = []
        for prim_data in primitives_payload:
            material_index = self._add_material(prim_data.get("material"))
            gltf_primitives.append(self._build_primitive(prim_data, material_index))
        mesh_index = len(self._meshes)
        self._meshes.append({
            "name": mesh_name,
            "primitives": gltf_primitives,
        })
        return mesh_index

    # ---------- materials and textures ----------

    def _add_material(self, mat):
        if not mat:
            mat = {
                "name": settings.DEFAULT_MATERIAL_NAME,
                "base_color": list(settings.DEFAULT_BASE_COLOR),
                "metallic": 0.0,
                "roughness": 1.0,
                "emissive": [0.0, 0.0, 0.0],
                "double_sided": False,
                "textures": {},
            }
        material_name = mat.get("name", settings.DEFAULT_MATERIAL_NAME)
        if material_name in self._material_cache_by_name:
            return self._material_cache_by_name[material_name]

        index = len(self._materials)
        uvt = mat.get("uv_transform")
        pbr = {
            "baseColorFactor": list(mat.get("base_color", [1.0, 1.0, 1.0, 1.0])),
            "metallicFactor": float(mat.get("metallic", 0.0)),
            "roughnessFactor": float(mat.get("roughness", 1.0)),
        }
        textures = mat.get("textures") or {}
        bc = textures.get("base_color")
        if bc:
            pbr["baseColorTexture"] = self._texture_info(bc, uvt)
        mr = textures.get("metallic_roughness")
        if mr:
            pbr["metallicRoughnessTexture"] = self._texture_info(mr, uvt)

        gltf_mat = {
            "name": mat.get("name", "material_{0}".format(index)),
            "pbrMetallicRoughness": pbr,
        }

        nrm = textures.get("normal")
        if nrm:
            extra = {}
            scale = mat.get("normal_scale")
            if scale is not None:
                try:
                    scale = float(scale)
                    if abs(scale - 1.0) > 1e-6:
                        extra["scale"] = scale
                except Exception:
                    pass
            gltf_mat["normalTexture"] = self._texture_info(nrm, uvt, extra)
        emi = textures.get("emissive")
        if emi:
            gltf_mat["emissiveTexture"] = self._texture_info(emi, uvt)
        occ = textures.get("occlusion")
        if occ:
            gltf_mat["occlusionTexture"] = self._texture_info(occ, uvt)

        emissive = mat.get("emissive") or [0.0, 0.0, 0.0]
        if any(abs(c) > 1e-9 for c in emissive):
            gltf_mat["emissiveFactor"] = [float(c) for c in emissive[:3]]
        if mat.get("double_sided"):
            gltf_mat["doubleSided"] = True

        # Glass / transmission. Blender imports both extensions.
        transmission = float(mat.get("transmission", 0.0) or 0.0)
        if transmission > 1e-4:
            mat_ext = gltf_mat.setdefault("extensions", {})
            mat_ext["KHR_materials_transmission"] = {"transmissionFactor": transmission}
            self._extensions_used.add("KHR_materials_transmission")
            ior = float(mat.get("ior", 1.5) or 1.5)
            if abs(ior - 1.5) > 1e-6:
                mat_ext["KHR_materials_ior"] = {"ior": ior}
                self._extensions_used.add("KHR_materials_ior")

        alpha_mode = mat.get("alpha_mode")
        if alpha_mode and alpha_mode != "OPAQUE":
            gltf_mat["alphaMode"] = alpha_mode
        if alpha_mode == "MASK":
            gltf_mat["alphaCutoff"] = float(mat.get("alpha_cutoff", 0.5))

        self._materials.append(gltf_mat)
        self._material_cache_by_name[material_name] = index
        return index

    def _texture_info(self, tex_info, uv_transform=None, extra=None):
        # Build a glTF textureInfo object, optionally with KHR_texture_transform
        # (UV tiling/offset). Blender's glTF importer reads this extension.
        info = dict(extra or {})
        info["index"] = self._add_texture(tex_info)
        if uv_transform:
            ext = {}
            offset = uv_transform.get("offset")
            scale = uv_transform.get("scale")
            if offset and (abs(offset[0]) > 1e-9 or abs(offset[1]) > 1e-9):
                ext["offset"] = [float(offset[0]), float(offset[1])]
            if scale and (abs(scale[0] - 1.0) > 1e-9 or abs(scale[1] - 1.0) > 1e-9):
                ext["scale"] = [float(scale[0]), float(scale[1])]
            if ext:
                info.setdefault("extensions", {})["KHR_texture_transform"] = ext
                self._extensions_used.add("KHR_texture_transform")
        return info

    def _add_texture(self, tex_info):
        # tex_info can be a string URI (external) or a dict with 'uri' or
        # 'embed_bytes' + 'mime_type'.
        if isinstance(tex_info, str):
            tex_info = {"uri": tex_info}
        if self._embed_textures and tex_info.get("embed_bytes") is not None:
            return self._add_embedded_texture(tex_info)
        uri = tex_info.get("uri", "")
        if uri in self._texture_cache_by_uri:
            return self._texture_cache_by_uri[uri]
        image = {"uri": _encode_uri(uri)}
        image_index = len(self._images)
        self._images.append(image)
        tex_index = len(self._textures)
        self._textures.append({"source": image_index})
        self._texture_cache_by_uri[uri] = tex_index
        return tex_index

    def _add_embedded_texture(self, tex_info):
        key = tex_info.get("source_path") or id(tex_info)
        if key in self._texture_cache_by_path:
            return self._texture_cache_by_path[key]
        data = tex_info["embed_bytes"]
        mime = tex_info.get("mime_type", "image/png")
        view_index = self._add_buffer_view(data)
        image_index = len(self._images)
        self._images.append({"mimeType": mime, "bufferView": view_index})
        tex_index = len(self._textures)
        self._textures.append({"source": image_index})
        self._texture_cache_by_path[key] = tex_index
        return tex_index

    # ---------- primitive assembly ----------

    def _build_primitive(self, prim_data, material_index):
        attributes = {}

        positions = prim_data.get("positions") or []
        if positions:
            attributes["POSITION"] = self._accessor_vec3_float(
                positions, target=settings.TARGET_ARRAY_BUFFER, with_min_max=True
            )

        normals = prim_data.get("normals") or []
        if normals:
            attributes["NORMAL"] = self._accessor_vec3_float(
                normals, target=settings.TARGET_ARRAY_BUFFER, with_min_max=False
            )

        tangents = prim_data.get("tangents") or []
        if tangents:
            attributes["TANGENT"] = self._accessor_vec4_float(
                tangents, target=settings.TARGET_ARRAY_BUFFER
            )

        uvs = prim_data.get("uvs") or []
        if uvs:
            attributes["TEXCOORD_0"] = self._accessor_vec2_float(
                uvs, target=settings.TARGET_ARRAY_BUFFER
            )

        colors = prim_data.get("colors") or []
        if colors:
            attributes["COLOR_0"] = self._accessor_vec4_float(
                colors, target=settings.TARGET_ARRAY_BUFFER
            )

        out = {"attributes": attributes, "material": material_index, "mode": 4}
        indices = prim_data.get("indices") or []
        if indices:
            out["indices"] = self._accessor_scalar_uint32(indices)
        return out

    # ---------- low-level buffer ----------

    def _align(self, alignment=settings.BUFFER_ALIGNMENT, fill=b"\x00"):
        r = len(self._buffer) % alignment
        if r:
            self._buffer.extend(fill * (alignment - r))

    def _add_buffer_view(self, data_bytes, target=None):
        self._align()
        offset = len(self._buffer)
        self._buffer.extend(data_bytes)
        view = {"buffer": 0, "byteOffset": offset, "byteLength": len(data_bytes)}
        if target is not None:
            view["target"] = target
        idx = len(self._buffer_views)
        self._buffer_views.append(view)
        return idx

    def _add_accessor(self, view_idx, component_type, count, type_str, mn=None, mx=None):
        acc = {
            "bufferView": view_idx,
            "componentType": component_type,
            "count": count,
            "type": type_str,
        }
        if mn is not None:
            acc["min"] = list(mn)
        if mx is not None:
            acc["max"] = list(mx)
        idx = len(self._accessors)
        self._accessors.append(acc)
        return idx

    def _accessor_vec3_float(self, flat, target, with_min_max=False):
        count = len(flat) // 3
        data = struct.pack("<{0}f".format(len(flat)), *flat)
        view = self._add_buffer_view(data, target=target)
        mn = mx = None
        if with_min_max and count:
            mn = [float("inf")] * 3
            mx = [float("-inf")] * 3
            for i in range(count):
                for k in range(3):
                    v = flat[i * 3 + k]
                    if v < mn[k]:
                        mn[k] = v
                    if v > mx[k]:
                        mx[k] = v
        return self._add_accessor(view, settings.COMPONENT_TYPE_FLOAT, count,
                                  settings.TYPE_VEC3, mn, mx)

    def _accessor_vec2_float(self, flat, target):
        count = len(flat) // 2
        data = struct.pack("<{0}f".format(len(flat)), *flat)
        view = self._add_buffer_view(data, target=target)
        return self._add_accessor(view, settings.COMPONENT_TYPE_FLOAT, count, settings.TYPE_VEC2)

    def _accessor_vec4_float(self, flat, target):
        count = len(flat) // 4
        data = struct.pack("<{0}f".format(len(flat)), *flat)
        view = self._add_buffer_view(data, target=target)
        return self._add_accessor(view, settings.COMPONENT_TYPE_FLOAT, count, settings.TYPE_VEC4)

    def _accessor_scalar_uint32(self, flat):
        count = len(flat)
        data = struct.pack("<{0}I".format(count), *flat)
        view = self._add_buffer_view(data, target=settings.TARGET_ELEMENT_ARRAY_BUFFER)
        return self._add_accessor(view, settings.COMPONENT_TYPE_UNSIGNED_INT, count, settings.TYPE_SCALAR)

    # ---------- output ----------

    def _build_root(self, bin_uri=None, buffer_byte_length=None):
        if buffer_byte_length is None:
            buffer_byte_length = len(self._buffer)
        root = {
            "asset": {
                "version": "2.0",
                "generator": settings.GENERATOR_STRING,
            },
            "scene": 0,
            "scenes": [{"nodes": list(self._scene_root_indices) or [0]}],
            "nodes": self._nodes,
            "meshes": self._meshes,
            "accessors": self._accessors,
            "bufferViews": self._buffer_views,
            "buffers": [{
                "byteLength": buffer_byte_length,
            }],
        }
        if bin_uri is not None:
            root["buffers"][0]["uri"] = _encode_uri(bin_uri)
        if self._materials:
            root["materials"] = self._materials
        if self._images:
            root["images"] = self._images
        if self._textures:
            root["textures"] = self._textures
        if self._samplers:
            root["samplers"] = self._samplers
        if self._cameras:
            root["cameras"] = self._cameras
        if self._lights:
            root.setdefault("extensions", {})["KHR_lights_punctual"] = {"lights": self._lights}
        if self._extensions_used:
            root["extensionsUsed"] = sorted(self._extensions_used)
        return root

    def write_gltf(self, output_path):
        out_dir = os.path.dirname(output_path)
        if out_dir and not os.path.isdir(out_dir):
            os.makedirs(out_dir)
        stem = os.path.splitext(os.path.basename(output_path))[0]
        bin_name = "{0}.bin".format(stem)
        bin_path = os.path.join(out_dir, bin_name) if out_dir else bin_name
        self._align()
        with open(bin_path, "wb") as f:
            f.write(self._buffer)
        gltf = self._build_root(bin_uri=bin_name)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(gltf, f, indent=2)
        return {"gltf_path": output_path, "bin_path": bin_path,
                "byte_length": len(self._buffer)}

    def write_glb(self, output_path):
        out_dir = os.path.dirname(output_path)
        if out_dir and not os.path.isdir(out_dir):
            os.makedirs(out_dir)
        # GLB layout: 12-byte header + JSON chunk + BIN chunk.
        self._align(4)
        bin_bytes = bytes(self._buffer)
        gltf = self._build_root(bin_uri=None, buffer_byte_length=len(bin_bytes))
        json_bytes = json.dumps(gltf, separators=(",", ":")).encode("utf-8")
        # JSON chunk padded with spaces to 4-byte boundary.
        json_pad = (4 - (len(json_bytes) % 4)) % 4
        json_bytes = json_bytes + (b" " * json_pad)
        # BIN chunk padded with zeros.
        bin_pad = (4 - (len(bin_bytes) % 4)) % 4
        bin_bytes = bin_bytes + (b"\x00" * bin_pad)

        total_length = 12 + 8 + len(json_bytes) + 8 + len(bin_bytes)
        with open(output_path, "wb") as f:
            f.write(struct.pack("<III", settings.GLB_MAGIC, settings.GLB_VERSION, total_length))
            f.write(struct.pack("<II", len(json_bytes), settings.GLB_CHUNK_JSON))
            f.write(json_bytes)
            f.write(struct.pack("<II", len(bin_bytes), settings.GLB_CHUNK_BIN))
            f.write(bin_bytes)
        return {"gltf_path": output_path, "bin_path": None,
                "byte_length": total_length}
