DigglesTool/internal/services/exporter/model.go

package exporter

import (
	"bytes"
	"fmt"
	"image"
	"image/png"
	"os"

	_ "github.com/ftrvxmtrx/tga"
	"github.com/qmuntal/gltf"
	"github.com/qmuntal/gltf/modeler"

	"git.influ.su/artmares/digglestool/internal/pkg/errors"
	"git.influ.su/artmares/digglestool/internal/pkg/logger"
	"git.influ.su/artmares/digglestool/pkg/threedb"
)

type Model struct {
	doc         *gltf.Document
	materialMap map[uint16]int
}

func (m *Model) Generate(model *threedb.Model, cache *Cache[*CacheItem], meshIndexes ...int) (*gltf.Document, error) {
	if model == nil {
		return nil, errors.ErrInvalidInput
	}

	m.doc = gltf.NewDocument()
	m.doc.Scenes = []*gltf.Scene{}
	m.materialMap = make(map[uint16]int)
	logger.Info("Generate use mesh indexes", len(meshIndexes))

	for i, meshIndex := range meshIndexes {
		if meshIndex >= len(model.Meshes) {
			return nil, errors.Wrapf(errors.ErrInvalidInput, "mesh index %d out of range (max: %d)",
				meshIndex, len(model.Meshes)-1)
		}

		mesh := model.Meshes[meshIndex]
		var nodeIndexes []int

		for index, link := range mesh.Links {
			materialIndex, err := m.materialIndex(model, cache, link.Material)
			if err != nil {
				return nil, errors.Wrapf(err, "failed to get material index for mesh %d link %d",
					meshIndex, index)
			}

			verticesIndex := m.generateVertices(model.Points[link.Points])
			textureCoordIndex := m.generateTextureCoordinate(model.TextureCoordinates[link.TextureCoordinates])
			trianglesIndex := modeler.WriteAccessor(m.doc, gltf.TargetArrayBuffer, model.Triangles[link.Triangles])

			m.doc.Meshes = append(m.doc.Meshes, &gltf.Mesh{
				Name: fmt.Sprintf("%s-%d-%d", model.Name, i, index),
				Primitives: []*gltf.Primitive{{
					Attributes: gltf.PrimitiveAttributes{
						gltf.POSITION:   verticesIndex,
						gltf.TEXCOORD_0: textureCoordIndex,
					},
					Indices:  gltf.Index(trianglesIndex),
					Material: gltf.Index(materialIndex),
				}},
			})

			nodeIndexes = append(nodeIndexes, len(m.doc.Nodes))
			m.doc.Nodes = append(m.doc.Nodes, &gltf.Node{
				Mesh: gltf.Index(len(m.doc.Meshes) - 1),
			})
		}

		m.doc.Scenes = append(m.doc.Scenes, &gltf.Scene{
			Name:  fmt.Sprintf("Root Scene-%d", meshIndex),
			Nodes: nodeIndexes,
		})
	}

	m.doc.Samplers = []*gltf.Sampler{{}}
	return m.doc, nil
}

func (m *Model) generateVertices(vectors []threedb.Vector) int {
	var vertices [][3]float32
	for _, vector := range vectors {
		vector = vector.Transform(100)
		vertices = append(vertices, vector.Normalize())
	}

	return modeler.WriteAccessor(m.doc, gltf.TargetElementArrayBuffer, vertices)
}

func (m *Model) generateTextureCoordinate(coordinates []threedb.Coordinate) int {
	var coords [][2]float32
	for _, coordinate := range coordinates {
		coords = append(coords, coordinate)
	}
	return modeler.WriteAccessor(m.doc, gltf.TargetArrayBuffer, coords)
}

func (m *Model) materialIndex(model *threedb.Model, cache *Cache[*CacheItem], meshMaterialIndex uint16) (int, error) {
	materialIndex, ok := m.materialMap[meshMaterialIndex]
	if !ok {
		material := model.Materials[meshMaterialIndex]
		item, ok := cache.Get(material.Name)
		if !ok {
			// Instead of returning nil, return a more descriptive error
			return 0, errors.New("material not found in cache: " + material.Name)
		}

		if item.Bytes == nil {
			f, err := os.Open(item.Path)
			if err != nil {
				return 0, errors.Wrapf(err, "failed to open texture file: %s", item.Path)
			}
			defer f.Close()

			img, _, err := image.Decode(f)
			if err != nil {
				return 0, errors.Wrap(err, "failed to decode image")
			}

			item.Bytes = &bytes.Buffer{}
			if err = png.Encode(item.Bytes, img); err != nil {
				return 0, errors.Wrap(err, "failed to encode image to PNG")
			}
			cache.Set(material.Name, item)
		}

		imageIndex, err := modeler.WriteImage(m.doc, material.Name, "image/png", item.Bytes)
		if err != nil {
			return 0, errors.Wrap(err, "failed to write image to GLTF document")
		}

		m.doc.Textures = append(m.doc.Textures, &gltf.Texture{
			Name:    material.Name,
			Sampler: gltf.Index(0),
			Source:  gltf.Index(imageIndex),
		})

		materialIndex = len(m.doc.Materials)
		m.doc.Materials = append(m.doc.Materials, &gltf.Material{
			Name:        material.Name,
			AlphaMode:   gltf.AlphaOpaque,
			AlphaCutoff: gltf.Float(0.5),
			PBRMetallicRoughness: &gltf.PBRMetallicRoughness{
				BaseColorTexture: &gltf.TextureInfo{Index: imageIndex},
				BaseColorFactor:  &[4]float64{0.8, 0.8, 0.8, 1},
				MetallicFactor:   gltf.Float(0.1),
				RoughnessFactor:  gltf.Float(0.99),
			},
		})
		m.materialMap[meshMaterialIndex] = materialIndex
	}

	return materialIndex, nil
}