// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// specialize is a low-level tool to generate type-specialized code. It is a
// convenience wrapper over text/template suitable for go generate. Unlike
// many other template tools, it does not parse Go code and allows use of
// text/template control within the template itself.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	noheader = flag.Bool("noheader", false, "Omit auto-generated header")
	pack     = flag.String("package", "", "Package name (optional)")
	imports  = flag.String("imports", "", "Comma-separated list of extra imports (optional)")

	x = flag.String("x", "", "Comma-separated list of X types (optional)")
	y = flag.String("y", "", "Comma-separated list of Y types (optional)")
	z = flag.String("z", "", "Comma-separated list of Z types (optional)")

	input  = flag.String("input", "", "Template file.")
	output = flag.String("output", "", "Filename for generated code. If not provided, a file next to the input is generated.")
)

// Top is the top-level struct to be passed to the template.
type Top struct {
	// Name is the base form of the filename: "foo/bar.tmpl" -> "bar".
	Name string
	// Package is the package name.
	Package string
	// Imports is a list of custom imports, if provided.
	Imports []string
	// X is the list of X type values.
	X []*X
}

// X is the concrete type to be iterated over in the user template.
type X struct {
	// Name is the name of X for use as identifier: "int" -> "Int", "[]byte" -> "ByteSlice".
	Name string
	// Type is the textual type of X: "int", "float32", "foo.Baz".
	Type string
	// Y is the list of Y type values for this X.
	Y []*Y
}

// Y is the concrete type to be iterated over in the user template for each X.
// Each combination of X and Y will be present.
type Y struct {
	// Name is the name of Y for use as identifier: "int" -> "Int", "[]byte" -> "ByteSlice".
	Name string
	// Type is the textual type of Y: "int", "float32", "foo.Baz".
	Type string
	// Z is the list of Z type values for this Y.
	Z []*Z
}

// Z is the concrete type to be iterated over in the user template for each Y.
// Each combination of X, Y and Z will be present.
type Z struct {
	// Name is the name of Z for use as identifier: "int" -> "Int", "[]byte" -> "ByteSlice".
	Name string
	// Type is the textual type of Z: "int", "float32", "foo.Baz".
	Type string
}

var (
	integers   = []string{"int", "int8", "int16", "int32", "int64", "uint", "uint8", "uint16", "uint32", "uint64"}
	floats     = []string{"float32", "float64"}
	primitives = append(append([]string{"bool", "string"}, integers...), floats...)

	macros = map[string][]string{
		"integers":   integers,
		"floats":     floats,
		"primitives": primitives,
		"data":       append([]string{"[]byte"}, primitives...),
		"universals": {"typex.T", "typex.U", "typex.V", "typex.W", "typex.X", "typex.Y", "typex.Z"},
	}

	packageMacros = map[string][]string{
		"typex": {"github.com/apache/beam/sdks/go/pkg/beam/core/typex"},
	}
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %v [options] --input=<filename.tmpl --x=<types>\n", filepath.Base(os.Args[0]))
	flag.PrintDefaults()
}

func main() {
	flag.Usage = usage
	flag.Parse()

	log.SetFlags(0)
	log.SetPrefix("specialize: ")

	if *input == "" {
		flag.Usage()
		log.Fatalf("no template file")
	}

	name := filepath.Base(*input)
	if index := strings.Index(name, "."); index > 0 {
		name = name[:index]
	}
	if *output == "" {
		*output = filepath.Join(filepath.Dir(*input), name+".go")
	}

	top := Top{Name: name, Package: *pack, Imports: expand(packageMacros, *imports)}
	var ys []*Y
	if *y != "" {
		var zs []*Z
		if *z != "" {
			for _, zt := range expand(macros, *z) {
				zs = append(zs, &Z{Name: makeName(zt), Type: zt})
			}
		}
		for _, yt := range expand(macros, *y) {
			ys = append(ys, &Y{Name: makeName(yt), Type: yt, Z: zs})
		}
	}
	for _, xt := range expand(macros, *x) {
		top.X = append(top.X, &X{Name: makeName(xt), Type: xt, Y: ys})
	}

	tmpl, err := template.New(*input).Funcs(funcMap).ParseFiles(*input)
	if err != nil {
		log.Fatalf("template parse failed: %v", err)
	}
	var buf bytes.Buffer
	if !*noheader {
		buf.WriteString("// File generated by specialize. Do not edit.\n\n")
	}
	if err := tmpl.Funcs(funcMap).Execute(&buf, top); err != nil {
		log.Fatalf("specialization failed: %v", err)
	}
	if err := ioutil.WriteFile(*output, buf.Bytes(), 0644); err != nil {
		log.Fatalf("write failed: %v", err)
	}
}

// expand parses, cleans up and expands macros for a comma-separated list.
func expand(subst map[string][]string, list string) []string {
	var ret []string
	for _, xt := range strings.Split(list, ",") {
		xt = strings.TrimSpace(xt)
		if xt == "" {
			continue
		}
		if exp, ok := subst[strings.ToLower(xt)]; ok {
			for _, t := range exp {
				ret = append(ret, t)
			}
			continue
		}
		ret = append(ret, xt)
	}
	return ret
}

// makeName creates a capitalized identifier from a type.
func makeName(t string) string {
	if strings.HasPrefix(t, "[]") {
		return makeName(t[2:] + "Slice")
	}

	t = strings.Replace(t, ".", "_", -1)
	t = strings.Replace(t, "[", "_", -1)
	t = strings.Replace(t, "]", "_", -1)
	return strings.Title(t)
}

// Useful template functions

var funcMap template.FuncMap = map[string]interface{}{
	"join":     strings.Join,
	"upto":     upto,
	"mkargs":   mkargs,
	"mktuple":  mktuple,
	"mktuplef": mktuplef,
}

// mkargs(n, type) returns "<fmt.Sprintf(format, 0)>, .., <fmt.Sprintf(format, n-1)> type".
// If n is 0, it returns the empty string.
func mkargs(n int, format, typ string) string {
	if n == 0 {
		return ""
	}
	return fmt.Sprintf("%v %v", mktuplef(n, format), typ)
}

// mktuple(n, v) returns "v, v, ..., v".
func mktuple(n int, v string) string {
	var ret []string
	for i := 0; i < n; i++ {
		ret = append(ret, v)
	}
	return strings.Join(ret, ", ")
}

// mktuplef(n, format) returns "<fmt.Sprintf(format, 0)>, .., <fmt.Sprintf(format, n-1)>"
func mktuplef(n int, format string) string {
	var ret []string
	for i := 0; i < n; i++ {
		ret = append(ret, fmt.Sprintf(format, i))
	}
	return strings.Join(ret, ", ")
}

// upto(n) returns []int{0, 1, .., n-1}.
func upto(i int) []int {
	var ret []int
	for k := 0; k < i; k++ {
		ret = append(ret, k)
	}
	return ret
}
