package main

// `mapart substances` turns the raw Fab/Quixel downloads into the handful of maps the landscape actually
// samples, at the resolution it actually needs.
//
// Why it exists. A downloaded set is nine maps at 4K - AO, BaseColor, Bump, Cavity, Displacement, Gloss,
// Normal, Roughness, Specular - about 110 MB a set, and the landscape material samples three of them. Worse,
// a UTexture2D keeps its *source* inside the uasset, so importing 4K would put hundreds of megabytes through
// LFS to render ground that is almost always seen at grazing distance: the scans are 2 m across, so 4K is
// 2048 pixels per metre. This writes 2K, three maps, and nothing else.
//
// Normals are downsampled and then **renormalised**. Averaging four unit vectors gives a shorter one, and a
// normal map whose vectors are not unit length lights slightly flat - not obviously wrong, just quietly
// duller everywhere, which is the kind of thing nobody finds later.
//
// The physical size comes out of the set's own metadata rather than being typed here. It is what the material
// needs to tile the texture life-size, and a number copied by hand is a number that goes stale when somebody
// swaps a substance for one scanned at a different scale.

import (
	"fmt"
	"image"
	"image/jpeg"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type substanceSet struct {
	Name      string  `json:"name"`
	Folder    string  `json:"folder"`
	Layer     string  `json:"layer"`
	PhysicalM float64 `json:"physical_m"` // 0 = read it from the set's metadata
	Note      string  `json:"note"`
}

type substanceConfig struct {
	SourceDir   string         `json:"source_dir"`
	OutputDir   string         `json:"output_dir"`
	Resolution  int            `json:"resolution"`
	Maps        []string       `json:"maps"`
	JpegQuality int            `json:"jpeg_quality"`
	Sets        []substanceSet `json:"sets"`
}

type groundManifest struct {
	Substances substanceConfig `json:"substances"`
}

// The per-asset metadata Quixel ships beside the maps. Only two fields matter here.
type quixelMeta struct {
	ID   string `json:"id"`
	Maps []struct {
		Name         string `json:"name"`
		Type         string `json:"type"`
		PhysicalSize string `json:"physicalSize"`
		Resolution   string `json:"resolution"`
		MimeType     string `json:"mimeType"`
	} `json:"maps"`
}

type substanceReport struct {
	Name      string   `json:"name"`
	Layer     string   `json:"layer"`
	Source    string   `json:"source"`
	PhysicalM float64  `json:"physical_m"`
	Written   []string `json:"written"`
	FromPx    int      `json:"from_px"`
	ToPx      int      `json:"to_px"`
	Seconds   float64  `json:"seconds"`
}

func substances(root string) error {
	groundPath := filepath.Join(root, "RawContent", "Terrain", "ground.json")
	var ground groundManifest
	if err := readJSON(groundPath, &ground); err != nil {
		return err
	}
	cfg := ground.Substances
	if len(cfg.Sets) == 0 {
		return fmt.Errorf("%s: no substances.sets, nothing to extract", groundPath)
	}
	if cfg.Resolution <= 0 {
		cfg.Resolution = 2048
	}
	if cfg.JpegQuality <= 0 {
		cfg.JpegQuality = 92
	}
	if len(cfg.Maps) == 0 {
		cfg.Maps = []string{"BaseColor", "Normal", "Roughness"}
	}
	srcRoot := filepath.Join(root, filepath.FromSlash(cfg.SourceDir))
	outDir := filepath.Join(root, filepath.FromSlash(cfg.OutputDir))
	if err := os.MkdirAll(outDir, 0o755); err != nil {
		return err
	}

	fmt.Printf("substances: %d set(s) -> %d px, maps %s\n", len(cfg.Sets), cfg.Resolution, strings.Join(cfg.Maps, ", "))

	reports := make([]substanceReport, 0, len(cfg.Sets))
	for _, set := range cfg.Sets {
		started := time.Now()
		dir := filepath.Join(srcRoot, filepath.FromSlash(set.Folder))
		extracted, err := findExtracted(dir)
		if err != nil {
			return fmt.Errorf("substance %q: %w", set.Name, err)
		}

		physical := set.PhysicalM
		if physical <= 0 {
			physical, err = physicalSize(extracted)
			if err != nil {
				return fmt.Errorf("substance %q: %w", set.Name, err)
			}
		}

		rep := substanceReport{Name: set.Name, Layer: set.Layer, Source: set.Folder, PhysicalM: physical, ToPx: cfg.Resolution}
		for _, kind := range cfg.Maps {
			srcFile, err := findMap(extracted, kind)
			if err != nil {
				return fmt.Errorf("substance %q: %w", set.Name, err)
			}
			im, err := load(srcFile)
			if err != nil {
				return fmt.Errorf("%s: %w", srcFile, err)
			}
			rep.FromPx = im.Bounds().Dx()

			isNormal := strings.EqualFold(kind, "Normal")
			out, err := downsampleSurface(im, cfg.Resolution, isNormal)
			if err != nil {
				return fmt.Errorf("%s: %w", srcFile, err)
			}

			// Normals as PNG: the source is already JPEG and a second lossy pass on a vector field shows up as
			// blocky lighting on flat ground. Colour and roughness stay JPEG, where it does not.
			name := fmt.Sprintf("T_%s_%s", set.Name, kind)
			var written string
			if isNormal {
				written = name + ".png"
				err = writePNG(filepath.Join(outDir, written), out)
			} else {
				written = name + ".jpg"
				err = writeJPEG(filepath.Join(outDir, written), out, cfg.JpegQuality)
			}
			if err != nil {
				return err
			}
			rep.Written = append(rep.Written, written)
		}
		rep.Seconds = time.Since(started).Seconds()
		fmt.Printf("  %-26s %-11s %4d -> %4d px, %.2f m scan, %d map(s), %.1fs\n",
			set.Name, set.Layer, rep.FromPx, rep.ToPx, rep.PhysicalM, len(rep.Written), rep.Seconds)
		reports = append(reports, rep)
	}

	if err := writeJSON(filepath.Join(outDir, "substances.json"), map[string]any{
		"when":       time.Now().UTC().Format(time.RFC3339),
		"resolution": cfg.Resolution,
		"maps":       cfg.Maps,
		"sets":       reports,
	}); err != nil {
		return err
	}
	fmt.Printf("%d substance(s) into %s\n", len(reports), outDir)
	return nil
}

// findExtracted locates the ..._extracted folder inside a Fab download, whose name nobody chose and which is
// the only place the maps and the metadata actually live.
func findExtracted(dir string) (string, error) {
	var found string
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() && strings.HasSuffix(info.Name(), "_extracted") {
			found = path
		}
		return nil
	})
	if err != nil {
		return "", err
	}
	if found == "" {
		return "", fmt.Errorf("no ..._extracted folder under %s - is this a Fab texture-set download?", dir)
	}
	return found, nil
}

func physicalSize(extracted string) (float64, error) {
	entries, err := os.ReadDir(extracted)
	if err != nil {
		return 0, err
	}
	for _, e := range entries {
		if !strings.EqualFold(filepath.Ext(e.Name()), ".json") {
			continue
		}
		var meta quixelMeta
		if err := readJSON(filepath.Join(extracted, e.Name()), &meta); err != nil {
			continue
		}
		for _, m := range meta.Maps {
			// "2x2" metres. Square scans only; a non-square one would need two numbers and the material
			// would need to know about both, so it is refused rather than silently halved.
			if m.PhysicalSize == "" {
				continue
			}
			parts := strings.Split(strings.ToLower(m.PhysicalSize), "x")
			if len(parts) != 2 {
				continue
			}
			w, err1 := strconv.ParseFloat(strings.TrimSpace(parts[0]), 64)
			h, err2 := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64)
			if err1 != nil || err2 != nil || w <= 0 {
				continue
			}
			if math.Abs(w-h) > 1e-6 {
				return 0, fmt.Errorf("scan is %s m, not square; the landscape material tiles with one number", m.PhysicalSize)
			}
			return w, nil
		}
	}
	return 0, fmt.Errorf("no physicalSize in the metadata under %s; set physical_m in ground.json instead", extracted)
}

func findMap(extracted, kind string) (string, error) {
	entries, err := os.ReadDir(extracted)
	if err != nil {
		return "", err
	}
	want := "_" + strings.ToLower(kind) + "."
	var best string
	for _, e := range entries {
		name := strings.ToLower(e.Name())
		if strings.Contains(name, want) && (strings.HasSuffix(name, ".jpg") || strings.HasSuffix(name, ".png")) {
			best = filepath.Join(extracted, e.Name())
		}
	}
	if best == "" {
		return "", fmt.Errorf("no %s map in %s", kind, extracted)
	}
	return best, nil
}

// downsampleSurface box-filters to `size` square. Colour is averaged in linear light like the map art;
// a normal map is averaged as a vector and renormalised, which is the whole reason this is not one function.
func downsampleSurface(im image.Image, size int, isNormal bool) (*image.RGBA, error) {
	read, err := rgbAccess(im)
	if err != nil {
		return nil, err
	}
	b := im.Bounds()
	w, h := b.Dx(), b.Dy()
	if size > w {
		size = w
	}
	sums := make([]float64, size*size*3)
	counts := make([]uint32, size*size)
	for y := 0; y < h; y++ {
		oy := y * size / h
		for x := 0; x < w; x++ {
			ox := x * size / w
			r, g, bl := read(b.Min.X+x, b.Min.Y+y)
			i := oy*size + ox
			if isNormal {
				// To -1..1 before averaging: the midpoint of two opposite normals is flat, and the midpoint
				// of their 0..1 encodings is flat too, but only if the average happens in the signed space.
				sums[i*3+0] += float64(r)/127.5 - 1
				sums[i*3+1] += float64(g)/127.5 - 1
				sums[i*3+2] += float64(bl)/127.5 - 1
			} else {
				sums[i*3+0] += float64(srgbToLinear[r])
				sums[i*3+1] += float64(srgbToLinear[g])
				sums[i*3+2] += float64(srgbToLinear[bl])
			}
			counts[i]++
		}
	}
	out := image.NewRGBA(image.Rect(0, 0, size, size))
	for i := 0; i < size*size; i++ {
		n := float64(counts[i])
		if n == 0 {
			n = 1
		}
		if isNormal {
			x, y, z := sums[i*3+0]/n, sums[i*3+1]/n, sums[i*3+2]/n
			length := math.Sqrt(x*x + y*y + z*z)
			if length < 1e-9 {
				x, y, z, length = 0, 0, 1, 1
			}
			x, y, z = x/length, y/length, z/length
			out.Pix[i*4+0] = encodeUnit(x)
			out.Pix[i*4+1] = encodeUnit(y)
			out.Pix[i*4+2] = encodeUnit(z)
		} else {
			out.Pix[i*4+0] = encodeSrgb(float32(sums[i*3+0] / n))
			out.Pix[i*4+1] = encodeSrgb(float32(sums[i*3+1] / n))
			out.Pix[i*4+2] = encodeSrgb(float32(sums[i*3+2] / n))
		}
		out.Pix[i*4+3] = 255
	}
	return out, nil
}

func encodeUnit(v float64) uint8 {
	return uint8(math.Round(math.Max(0, math.Min(255, (v+1)*127.5))))
}

func writeJPEG(path string, im image.Image, quality int) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	if err := jpeg.Encode(f, im, &jpeg.Options{Quality: quality}); err != nil {
		f.Close()
		return err
	}
	return f.Close()
}