package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"text/template"
)

type opType struct {
	Op   string
	Code string
}

// fieldRunTypes are the kinds of a field which have the opcodes of a run: see internal/encoder.fieldRunOps.
var fieldRunTypes = []string{"Int", "Uint", "Float64", "String", "Bool"}

func createOpType(op, code string) opType {
	return opType{
		Op:   op,
		Code: code,
	}
}

func _main() error {
	tmpl, err := template.New("").Parse(`// Code generated by internal/cmd/generator. DO NOT EDIT.
package encoder

import (
  "strings"
)

type CodeType int

const (
{{- range $index, $type := .CodeTypes }}
  Code{{ $type }} CodeType = {{ $index }}
{{- end }}
)

var opTypeStrings = [{{ .OpLen }}]string{
{{- range $type := .OpTypes }}
    "{{ $type.Op }}",
{{- end }}
}

type OpType uint16

const (
{{- range $index, $type := .OpTypes }}
  Op{{ $type.Op }} OpType = {{ $index }}
{{- end }}
)

func (t OpType) String() string {
  if int(t) >= {{ .OpLen }} {
    return ""
  }
  return opTypeStrings[int(t)]
}

func (t OpType) CodeType() CodeType {
  if strings.Contains(t.String(), "Struct") {
    if strings.Contains(t.String(), "End") {
      return CodeStructEnd
    }
    return CodeStructField
  }
  switch t {
  case OpArray, OpArrayPtr:
    return CodeArrayHead
  case OpArrayElem:
    return CodeArrayElem
  case OpSlice, OpSlicePtr:
    return CodeSliceHead
  case OpSliceElem:
    return CodeSliceElem
  case OpMap, OpMapPtr:
    return CodeMapHead
  case OpMapKey:
    return CodeMapKey
  case OpMapValue:
    return CodeMapValue
  case OpMapEnd:
    return CodeMapEnd
  }

  return CodeOp
}

func (t OpType) HeadToPtrHead() OpType {
  if t == OpStructHead {
    return OpStructPtrHead
  }
  return t
}

func (t OpType) PtrHeadToHead() OpType {
  if t == OpStructPtrHead {
    return OpStructHead
  }
  return t
}

func (t OpType) FieldToEnd() OpType {
  idx := strings.Index(t.String(), "Field")
  if idx == -1 {
    return t
  }
  suffix := t.String()[idx+len("Field"):]
  if suffix == "" || suffix == "OmitEmpty" {
    return t
  }
  const toEndOffset = 2
  if strings.Contains(OpType(int(t) + toEndOffset).String(), "End"+suffix) {
    return OpType(int(t) + toEndOffset)
  }
  return t
}

func (t OpType) FieldToOmitEmptyField() OpType {
  const toOmitEmptyOffset = 1
  if strings.Contains(OpType(int(t) + toOmitEmptyOffset).String(), "OmitEmpty") {
    return OpType(int(t) + toOmitEmptyOffset)
  }
  return t
}
`)
	if err != nil {
		return err
	}
	codeTypes := []string{
		"Op",
		"ArrayHead",
		"ArrayElem",
		"SliceHead",
		"SliceElem",
		"MapHead",
		"MapKey",
		"MapValue",
		"MapEnd",
		"Recursive",
		"StructField",
		"StructEnd",
	}
	primitiveTypes := []string{
		"int", "uint", "float32", "float64", "bool", "string", "bytes", "number",
		"array", "map", "slice", "struct", "MarshalJSON", "MarshalText",
		"intString", "uintString", "float32String", "float64String", "boolString", "stringString", "numberString",
		"intPtr", "uintPtr", "float32Ptr", "float64Ptr", "boolPtr", "stringPtr", "bytesPtr", "numberPtr",
		"arrayPtr", "mapPtr", "slicePtr", "marshalJSONPtr", "marshalTextPtr", "interfacePtr",
		"intPtrString", "uintPtrString", "float32PtrString", "float64PtrString", "boolPtrString", "stringPtrString", "numberPtrString",
	}
	primitiveTypesUpper := []string{}
	for _, typ := range primitiveTypes {
		primitiveTypesUpper = append(primitiveTypesUpper, strings.ToUpper(string(typ[0]))+typ[1:])
	}
	opTypes := []opType{
		createOpType("End", "Op"),
		createOpType("Interface", "Op"),
		createOpType("Ptr", "Op"),
		createOpType("SliceElem", "SliceElem"),
		createOpType("SliceEnd", "Op"),
		createOpType("ArrayElem", "ArrayElem"),
		createOpType("ArrayEnd", "Op"),
		createOpType("MapKey", "MapKey"),
		createOpType("MapValue", "MapValue"),
		createOpType("MapEnd", "Op"),
		createOpType("Recursive", "Op"),
		createOpType("RecursivePtr", "Op"),
		createOpType("RecursiveEnd", "Op"),
		createOpType("InterfaceEnd", "Op"),
	}
	for _, typ := range primitiveTypesUpper {
		opTypes = append(opTypes, createOpType(typ, "Op"))
	}
	// the head of a struct is an opcode of its own: it is followed by the opcode of the first field.
	opTypes = append(opTypes,
		opType{Op: "StructHead", Code: "StructField"},
		opType{Op: "StructPtrHead", Code: "StructField"},
	)
	// a field of interface{} has the opcodes of a field of its own, which encode the value as OpInterface does:
	// the opcode of a value of interface{} is among the first ones, so it is not a primitive type here.
	for _, typ := range append(primitiveTypesUpper, "Interface", "") {
		for _, opt := range []string{"", "OmitEmpty"} {
			op := fmt.Sprintf(
				"StructField%s%s",
				opt,
				typ,
			)
			opTypes = append(opTypes, opType{
				Op:   op,
				Code: "StructField",
			})
		}
		for _, opt := range []string{"", "OmitEmpty"} {
			op := fmt.Sprintf(
				"StructEnd%s%s",
				opt,
				typ,
			)
			opTypes = append(opTypes, opType{
				Op:   op,
				Code: "StructEnd",
			})
		}
	}
	// A run of fields of the same kind in a row is encoded by the opcode of the run, which falls through the
	// cases of the shorter runs: the fields after the first are encoded without a dispatch. The opcodes come
	// after every other one, so that the offsets between the opcodes of a field, its omitempty and its end
	// stay as they are.
	for _, typ := range fieldRunTypes {
		for _, n := range []string{"3", "2"} {
			opTypes = append(opTypes, opType{Op: "StructField" + typ + n, Code: "StructField"})
		}
	}
	var b bytes.Buffer
	if err := tmpl.Execute(&b, struct {
		CodeTypes []string
		OpTypes   []opType
		OpLen     int
	}{
		CodeTypes: codeTypes,
		OpTypes:   opTypes,
		OpLen:     len(opTypes),
	}); err != nil {
		return err
	}
	path := filepath.Join(repoRoot(), "internal", "encoder", "optype.go")
	buf, err := format.Source(b.Bytes())
	if err != nil {
		return err
	}
	return os.WriteFile(path, buf, 0644)
}

// encoderFunctions are the functions of the template which the plain VM has as they are in the encoder
// package: the other VMs define them with their indent or their colors.
var encoderFunctions = map[string]string{
	"appendInt":       "AppendInt",
	"appendUint":      "AppendUint",
	"appendFloat64":   "AppendFloat64",
	"appendString":    "AppendString",
	"appendByteSlice": "AppendByteSlice",
	"appendNumber":    "AppendNumber",
}

// callEncoderDirectly makes the calls of encoderFunctions calls of the encoder package: a call through a
// function variable is a load and an indirect call, for every value the VM appends.
func callEncoderDirectly(f *ast.File) {
	ast.Inspect(f, func(n ast.Node) bool {
		call, ok := n.(*ast.CallExpr)
		if !ok {
			return true
		}
		ident, ok := call.Fun.(*ast.Ident)
		if !ok {
			return true
		}
		if name, ok := encoderFunctions[ident.Name]; ok {
			call.Fun = &ast.SelectorExpr{X: ast.NewIdent("encoder"), Sel: ast.NewIdent(name)}
		}
		return true
	})
}

func generateVM() error {
	file, err := os.ReadFile("vm.go.tmpl")
	if err != nil {
		return err
	}
	for _, pkg := range []string{"vm", "vm_indent", "vm_color", "vm_color_indent"} {
		// parsed for each package: the tree of the plain VM is changed.
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, "", string(file), parser.ParseComments)
		if err != nil {
			return err
		}
		f.Name.Name = pkg
		if pkg == "vm" {
			callEncoderDirectly(f)
		}
		var buf bytes.Buffer
		printer.Fprint(&buf, fset, f)
		path := filepath.Join(repoRoot(), "internal", "encoder", pkg, "vm.go")
		source, err := format.Source(buf.Bytes())
		if err != nil {
			return err
		}
		if err := os.WriteFile(path, source, 0644); err != nil {
			return err
		}
	}
	return nil
}

func repoRoot() string {
	_, file, _, _ := runtime.Caller(0)
	relativePathFromRepoRoot := filepath.Join("internal", "cmd", "generator")
	return strings.TrimSuffix(filepath.Dir(file), relativePathFromRepoRoot)
}

//go:generate go run main.go
func main() {
	if err := generateVM(); err != nil {
		panic(err)
	}
	if err := _main(); err != nil {
		panic(err)
	}
}