// SPDX-License-Identifier: Unlicense OR MIT

// +build ignore

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"text/template"
)

// This program generates shader variants for
// multiple GPU backends (OpenGL ES, Direct3D 11...)
// from a single source.

type shaderArgs struct {
	FetchColorExpr string
	Header         string
}

// TextureBinding matches gpu.TextureBinding.
type TextureBinding struct {
	Name    string
	Binding int
}

// InputLocation matches gpu.InputLocation.
type InputLocation struct {
	Name          string
	Location      int
	Semantic      string
	SemanticIndex int

	Type DataType
	Size int
}

type DataType uint8

// UniformLocation matches gpu.UniformLocation.
type UniformLocation struct {
	Name   string
	Type   DataType
	Size   int
	Offset int
}

const (
	DataTypeFloat DataType = iota
	DataTypeShort
)

func main() {
	if err := generate(); err != nil {
		fmt.Fprintf(os.Stderr, "gpu generate: %v\n", err)
		os.Exit(1)
	}
}

func generate() error {
	tmp, err := ioutil.TempDir("", "gpu-generate")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmp)
	glslcc, err := exec.LookPath("glslcc")
	if err != nil {
		return err
	}
	fxc, err := exec.LookPath("fxc")
	fxcFound := err == nil
	shaders, err := filepath.Glob("shaders/*")
	if err != nil {
		return err
	}
	var out bytes.Buffer
	out.WriteString("// Code generated by build.go. DO NOT EDIT.\n\n")
	out.WriteString("package gpu\n\n")

	out.WriteString("var (\n")

	for _, shader := range shaders {
		const nvariants = 2
		var variants [nvariants]struct {
			gles2       string
			hlslSrc     string
			hlsl        []byte
			inputs      []InputLocation
			uniforms    []UniformLocation
			textures    []TextureBinding
			uniformSize int
		}
		args := [nvariants]shaderArgs{
			{
				FetchColorExpr: `_color`,
				Header:         `layout(binding=0) uniform Color { vec4 _color; };`,
			},
			{
				FetchColorExpr: `texture(tex, vUV)`,
				Header:         `layout(binding=0) uniform sampler2D tex;`,
			},
		}
		for i := range args {
			gles2, reflect, err := convertShader(tmp, glslcc, shader, "gles", "100", &args[i], false)
			if err != nil {
				return err
			}
			// Make the GL ES 2 source compatible with desktop GL 3.
			gles2 = "#version 100\n" + gles2
			inputs, uniforms, textures, uniformSize, err := parseReflection(reflect)
			if err != nil {
				return err
			}
			hlsl, _, err := convertShader(tmp, glslcc, shader, "hlsl", "40", &args[i], false)
			if err != nil {
				return err
			}
			var hlslProf string
			switch filepath.Ext(shader) {
			case ".frag":
				hlslProf = "ps"
			case ".vert":
				hlslProf = "vs"
			default:
				return fmt.Errorf("unrecognized shader type %s", shader)
			}
			var hlslc []byte
			if fxcFound {
				hlslc, err = compileHLSL(tmp, fxc, hlsl, "main", hlslProf+"_4_0")
				if err != nil {
					return err
				}
			}
			variants[i].gles2 = gles2
			variants[i].hlslSrc = hlsl
			variants[i].hlsl = hlslc
			variants[i].inputs = inputs
			variants[i].uniforms = uniforms
			variants[i].textures = textures
			variants[i].uniformSize = uniformSize
		}
		name := filepath.Base(shader)
		name = strings.ReplaceAll(name, ".", "_")
		fmt.Fprintf(&out, "\tshader_%s = ", name)
		// If the shader don't use the variant arguments, output
		// only a single version.
		multiVariant := variants[0].gles2 != variants[1].gles2
		if multiVariant {
			fmt.Fprintf(&out, "[...]ShaderSources{\n")
		}
		for _, src := range variants {
			fmt.Fprintf(&out, "ShaderSources{\n")
			if len(src.inputs) > 0 {
				fmt.Fprintf(&out, "Inputs: []InputLocation{\n")
				for _, inp := range src.inputs {
					fmt.Fprintf(&out, "{Name: %q, Location: %d, Semantic: %q, ", inp.Name, inp.Location, inp.Semantic)
					fmt.Fprintf(&out, "SemanticIndex: %d, Type: %d, Size: %d},\n", inp.SemanticIndex, inp.Type, inp.Size)
				}
				fmt.Fprintf(&out, "},\n")
			}
			if len(src.uniforms) > 0 {
				fmt.Fprintf(&out, "Uniforms: []UniformLocation{\n")
				for _, u := range src.uniforms {
					fmt.Fprintf(&out, "{Name: %q, Type: %d, Size: %d, Offset: %d},\n", u.Name, u.Type, u.Size, u.Offset)
				}
				fmt.Fprintf(&out, "},\n")
			}
			if src.uniformSize != 0 {
				fmt.Fprintf(&out, "UniformSize: %d,\n", src.uniformSize)
			}
			if len(src.textures) > 0 {
				fmt.Fprintf(&out, "Textures: []TextureBinding{\n")
				for _, t := range src.textures {
					fmt.Fprintf(&out, "{Name: %q, Binding: %d},\n", t.Name, t.Binding)
				}
				fmt.Fprintf(&out, "},\n")
			}
			fmt.Fprintf(&out, "GLES2: %#v,\n", src.gles2)
			fmt.Fprintf(&out, "/*\n%s\n*/\n", src.hlslSrc)
			fmt.Fprintf(&out, "HLSL: %#v,\n", src.hlsl)
			fmt.Fprintf(&out, "}")
			if multiVariant {
				fmt.Fprintf(&out, ",")
			}
			fmt.Fprintf(&out, "\n")
			if !multiVariant {
				break
			}
		}
		if multiVariant {
			fmt.Fprintf(&out, "}\n")
		}
	}
	out.WriteString(")")
	gosrc, err := format.Source(out.Bytes())
	if err != nil {
		return fmt.Errorf("shader.go: %v", err)
	}
	return ioutil.WriteFile("shaders.go", gosrc, 0644)
}

func parseReflection(jsonData []byte) ([]InputLocation, []UniformLocation, []TextureBinding, int, error) {
	type InputReflection struct {
		ID            int    `json:"id"`
		Name          string `json:"name"`
		Location      int    `json:"location"`
		Semantic      string `json:"semantic"`
		SemanticIndex int    `json:"semantic_index"`
		Type          string `json:"type"`
	}
	type UniformMemberReflection struct {
		Name   string `json:"name"`
		Type   string `json:"type"`
		Offset int    `json:"offset"`
		Size   int    `json:"size"`
	}
	type UniformBufferReflection struct {
		ID      int                       `json:"id"`
		Name    string                    `json:"name"`
		Set     int                       `json:"set"`
		Binding int                       `json:"binding"`
		Size    int                       `json:"block_size"`
		Members []UniformMemberReflection `json:"members"`
	}
	type TextureReflection struct {
		ID        int    `json:"id"`
		Name      string `json:"name"`
		Set       int    `json:"set"`
		Binding   int    `json:"binding"`
		Dimension string `json:"dimension"`
		Format    string `json:"format"`
	}
	type shaderReflection struct {
		Inputs         []InputReflection         `json:"inputs"`
		UniformBuffers []UniformBufferReflection `json:"uniform_buffers"`
		Textures       []TextureReflection       `json:"textures"`
	}
	type shaderMetadata struct {
		VS shaderReflection `json:"vs"`
		FS shaderReflection `json:"fs"`
	}
	var reflect shaderMetadata
	if err := json.Unmarshal(jsonData, &reflect); err != nil {
		return nil, nil, nil, 0, fmt.Errorf("parseReflection: %v", err)
	}
	var inputs []InputLocation
	inputRef := reflect.VS.Inputs
	for _, input := range inputRef {
		dataType, dataSize, err := parseDataType(input.Type)
		if err != nil {
			return nil, nil, nil, 0, fmt.Errorf("parseReflection: %v", err)
		}
		inputs = append(inputs, InputLocation{
			Name:          input.Name,
			Location:      input.Location,
			Semantic:      input.Semantic,
			SemanticIndex: input.SemanticIndex,
			Type:          dataType,
			Size:          dataSize,
		})
	}
	sort.Slice(inputs, func(i, j int) bool {
		return inputs[i].Location < inputs[j].Location
	})
	var ublocks []UniformLocation
	shaderBlocks := reflect.VS.UniformBuffers
	if len(shaderBlocks) == 0 {
		shaderBlocks = reflect.FS.UniformBuffers
	}
	blockOffset := 0
	for _, block := range shaderBlocks {
		for _, member := range block.Members {
			dataType, size, err := parseDataType(member.Type)
			if err != nil {
				return nil, nil, nil, 0, fmt.Errorf("parseReflection: %v", err)
			}
			ublocks = append(ublocks, UniformLocation{
				// Synthetic name generated by glslcc.
				Name:   fmt.Sprintf("_%d.%s", block.ID, member.Name),
				Type:   dataType,
				Size:   size,
				Offset: blockOffset + member.Offset,
			})
		}
		blockOffset += block.Size
	}
	textures := reflect.VS.Textures
	if len(textures) == 0 {
		textures = reflect.FS.Textures
	}
	var texBinds []TextureBinding
	for _, texture := range textures {
		texBinds = append(texBinds, TextureBinding{
			Name:    texture.Name,
			Binding: texture.Binding,
		})
	}
	return inputs, ublocks, texBinds, blockOffset, nil
}

func parseDataType(t string) (DataType, int, error) {
	switch t {
	case "float":
		return DataTypeFloat, 1, nil
	case "float2":
		return DataTypeFloat, 2, nil
	case "float3":
		return DataTypeFloat, 3, nil
	case "float4":
		return DataTypeFloat, 4, nil
	default:
		return 0, 0, fmt.Errorf("unsupported input data type: %s", t)
	}
}

func compileHLSL(tmp, fxc, src, entry, profile string) ([]byte, error) {
	tmpfile := filepath.Join(tmp, "shader.hlsl")
	if err := ioutil.WriteFile(tmpfile, []byte(src), 0644); err != nil {
		return nil, err
	}
	outFile := filepath.Join(tmp, "shader.bin")
	cmd := exec.Command(fxc,
		"/T", profile,
		"/E", entry,
		"/nologo",
		"/Fo", outFile,
		tmpfile,
	)
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	return ioutil.ReadFile(outFile)
}

func convertShader(tmp, glslcc, path, lang, profile string, args *shaderArgs, flattenUBOs bool) (string, []byte, error) {
	shaderTmpl, err := template.ParseFiles(path)
	if err != nil {
		return "", nil, err
	}
	var buf bytes.Buffer
	if err := shaderTmpl.Execute(&buf, args); err != nil {
		return "", nil, err
	}
	tmppath := filepath.Join(tmp, filepath.Base(path))
	if err := ioutil.WriteFile(tmppath, buf.Bytes(), 0644); err != nil {
		return "", nil, err
	}
	defer os.Remove(tmppath)
	var progFlag string
	var progSuffix string
	switch filepath.Ext(path) {
	case ".vert":
		progFlag = "--vert"
		progSuffix = "vs"
	case ".frag":
		progFlag = "--frag"
		progSuffix = "fs"
	default:
		return "", nil, fmt.Errorf("unrecognized shader type: %s", path)
	}
	cmd := exec.Command(glslcc,
		"--silent",
		"--optimize",
		"--reflect",
		"--output", filepath.Join(tmp, "shader"),
		"--lang", lang,
		"--profile", profile,
		progFlag, tmppath,
	)
	if flattenUBOs {
		cmd.Args = append(cmd.Args, "--flatten-ubos")
	}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return "", nil, fmt.Errorf("%s: %v", path, err)
	}
	f, err := os.Open(filepath.Join(tmp, "shader_"+progSuffix))
	if err != nil {
		return "", nil, err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	src, err := ioutil.ReadAll(f)
	if err != nil {
		return "", nil, err
	}
	reflect, err := ioutil.ReadFile(filepath.Join(tmp, "shader_"+progSuffix+".json"))
	if err != nil {
		return "", nil, err
	}
	return string(src), reflect, nil
}