// Copyright 2018 The CUE Authors
//
// Licensed 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.

//go:build ignore

// gen.go generates the pkg.go files inside the packages under the pkg directory.
//
// It takes the list of packages from the packages.txt.
//
// Be sure to also update an entry in pkg/pkg.go, if so desired.
package main

import (
	"bytes"
	"cmp"
	_ "embed"
	"flag"
	"fmt"
	"go/constant"
	"go/format"
	"go/token"
	"go/types"
	"log"
	"math/big"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"text/template"

	"golang.org/x/tools/go/packages"

	"cuelang.org/go/cue/ast"
	cueformat "cuelang.org/go/cue/format"
	"cuelang.org/go/cue/parser"
	"cuelang.org/go/internal"
)

type headerParams struct {
	GoPkg  string
	CUEPkg string

	PackageDoc  string
	PackageDefs string
}

var header = template.Must(template.New("").Parse(
	`// Code generated by cuelang.org/go/pkg/gen. DO NOT EDIT.

{{if .PackageDoc}}
{{.PackageDoc -}}
//     {{.PackageDefs}}
{{end -}}
package {{.GoPkg}}

{{if .CUEPkg -}}
import (
	"cuelang.org/go/internal/core/adt"
	"cuelang.org/go/internal/pkg"
)

func init() {
	pkg.Register({{printf "%q" .CUEPkg}}, p)
}

var _ = adt.TopKind // in case the adt package isn't used
{{end}}
`))

const pkgParent = "cuelang.org/go/pkg"

func main() {
	flag.Parse()
	log.SetFlags(log.Lshortfile)
	log.SetOutput(os.Stdout)

	cfg := &packages.Config{Mode: packages.NeedName | packages.NeedFiles | packages.NeedTypes}
	pkgs, err := packages.Load(cfg, "./...")
	if err != nil {
		fmt.Fprintf(os.Stderr, "load: %v\n", err)
		os.Exit(1)
	}
	if packages.PrintErrors(pkgs) > 0 {
		os.Exit(1)
	}
	// Sort the Go packages by import path; otherwise adding a new builtin package
	// puts it at the very end of the list the first time it is getting added to register.go,
	// as it's not imported by the root package yet. Sorting ensures consistent output.
	slices.SortFunc(pkgs, func(a, b *packages.Package) int {
		return cmp.Compare(a.PkgPath, b.PkgPath)
	})

	regBuf := new(bytes.Buffer)
	fmt.Fprintf(regBuf, "// Code generated by cuelang.org/go/pkg/gen. DO NOT EDIT.\n\n")
	fmt.Fprintf(regBuf, "package pkg\n\n")
	fmt.Fprintf(regBuf, "import (\n")
	for _, pkg := range pkgs {
		switch {
		case pkg.PkgPath == pkgParent:
			// The pkg package itself should not be generated.
		case strings.Contains(pkg.PkgPath, "/internal"):
			// Internal packages are not for public use.
		default:
			fmt.Fprintf(regBuf, "\t_ %q\n", pkg.PkgPath)
			if pkg.PkgPath == "cuelang.org/go/pkg/path" {
				// TODO remove this special case. Currently the path
				// pkg.go file cannot be generated automatically but that
				// will be possible when we can attach arbitrary signatures
				// to builtin functions.
				break
			}
			if err := generate(pkg); err != nil {
				log.Fatalf("%s: %v", pkg, err)
			}
		}
	}
	fmt.Fprintf(regBuf, ")\n")
	if err := os.WriteFile("register.go", regBuf.Bytes(), 0o666); err != nil {
		log.Fatal(err)
	}
}

type generator struct {
	dir         string
	w           *bytes.Buffer
	cuePkgPath  string
	first       bool
	nonConcrete bool
}

func generate(pkg *packages.Package) error {
	// go/packages supports multiple build systems, including some which don't keep
	// a Go package entirely within a single directory.
	// However, we know for certain that CUE uses modules, so it is the case here.
	// We can figure out the directory from the first Go file.
	pkgDir := filepath.Dir(pkg.GoFiles[0])
	cuePkg := strings.TrimPrefix(pkg.PkgPath, pkgParent+"/")
	g := generator{
		dir:        pkgDir,
		cuePkgPath: cuePkg,
		w:          &bytes.Buffer{},
	}

	params := headerParams{
		GoPkg:  pkg.Name,
		CUEPkg: cuePkg,
	}
	// As a special case, the "tool" package cannot be imported from CUE.
	skipRegister := params.CUEPkg == "tool"
	if skipRegister {
		params.CUEPkg = ""
	}

	if doc, err := os.ReadFile(filepath.Join(pkgDir, "doc.txt")); err == nil {
		defs, err := os.ReadFile(filepath.Join(pkgDir, pkg.Name+".cue"))
		if err != nil {
			return err
		}
		i := bytes.Index(defs, []byte("package "+pkg.Name))
		defs = defs[i+len("package "+pkg.Name)+1:]
		defs = bytes.TrimRight(defs, "\n")
		defs = bytes.ReplaceAll(defs, []byte("\n"), []byte("\n//\t"))
		params.PackageDoc = string(doc)
		params.PackageDefs = string(defs)
	}

	if err := header.Execute(g.w, params); err != nil {
		return err
	}

	if !skipRegister {
		fmt.Fprintf(g.w, "var p = &pkg.Package{\nNative: []*pkg.Builtin{")
		g.first = true
		if err := g.processGo(pkg); err != nil {
			return err
		}
		fmt.Fprintf(g.w, "},\n")
		if err := g.processCUE(); err != nil {
			return err
		}
		fmt.Fprintf(g.w, "}\n")
	}

	b, err := format.Source(g.w.Bytes())
	if err != nil {
		fmt.Printf("go/format error on %s: %v\n", pkg.PkgPath, err)
		b = g.w.Bytes() // write the unformatted source
	}

	filename := filepath.Join(pkgDir, "pkg.go")

	if err := os.WriteFile(filename, b, 0666); err != nil {
		return err
	}
	return nil
}

func (g *generator) sep() {
	if g.first {
		g.first = false
		return
	}
	fmt.Fprint(g.w, ", ")
}

// processCUE mixes in CUE definitions defined in the package directory.
func (g *generator) processCUE() error {
	// Note: we avoid using the cue/load and the cuecontext packages
	// because they depend on the standard library which is what this
	// command is generating - cyclic dependencies are undesirable in general.
	// We only need to load the declarations from one CUE file if it exists.
	expr, err := loadCUEDecls(g.dir)
	if err != nil {
		return fmt.Errorf("error processing %s: %v", g.cuePkgPath, err)
	}
	if expr == nil { // No syntax to add.
		return nil
	}
	b, err := cueformat.Node(expr)
	if err != nil {
		return err
	}

	// Compact the CUE by removing empty lines. This requires re-formatting to align fields.
	// TODO(mvdan): provide a "compact" option in cue/format for this purpose?
	b = bytes.ReplaceAll(b, []byte("\n\n"), []byte("\n"))
	b, err = cueformat.Source(b)
	if err != nil {
		return err
	}
	b = bytes.TrimSpace(b) // no trailing newline

	// Try to use a Go string with backquotes, for readability.
	// If not possible due to cueSrc itself having backquotes,
	// use a single-line double-quoted string, removing tabs for brevity.
	// We don't use strconv.CanBackquote as it is for quoting as a single line.
	if cueSrc := string(b); !strings.Contains(cueSrc, "`") {
		fmt.Fprintf(g.w, "CUE: `%s`,\n", cueSrc)
	} else {
		cueSrc = strings.ReplaceAll(cueSrc, "\t", "")
		fmt.Fprintf(g.w, "CUE: %q,\n", cueSrc)
	}
	return nil
}

func (g *generator) processGo(pkg *packages.Package) error {
	// We sort the objects by their original source code position.
	// Otherwise, go/types defaults to sorting by name strings.
	// We could remove this code if we were fine with sorting by name.
	scope := pkg.Types.Scope()
	type objWithPos struct {
		obj types.Object
		pos token.Position
	}
	var objs []objWithPos
	for _, name := range scope.Names() {
		obj := scope.Lookup(name)
		objs = append(objs, objWithPos{obj, pkg.Fset.Position(obj.Pos())})
	}
	slices.SortFunc(objs, func(a, b objWithPos) int {
		if c := cmp.Compare(a.pos.Filename, b.pos.Filename); c != 0 {
			return c
		}
		return cmp.Compare(a.pos.Line, b.pos.Line)
	})

	for _, obj := range objs {
		obj := obj.obj // no longer need the token.Position
		if !obj.Exported() {
			continue
		}
		// TODO: support type declarations.
		switch obj := obj.(type) {
		case *types.Const:
			var value string
			switch v := obj.Val(); v.Kind() {
			case constant.Bool, constant.Int, constant.String:
				// TODO: convert octal numbers
				value = v.ExactString()
			case constant.Float:
				var rat big.Rat
				rat.SetString(v.ExactString())
				var float big.Float
				float.SetRat(&rat)
				value = float.Text('g', -1)
			default:
				fmt.Printf("Dropped entry %s.%s (%T: %v)\n", g.cuePkgPath, obj.Name(), v.Kind(), v.ExactString())
				continue
			}
			g.sep()
			fmt.Fprintf(g.w, "{\nName: %q,\n Const: %q,\n}", obj.Name(), value)
		case *types.Func:
			g.genFunc(obj)
		}
	}
	return nil
}

var (
	typeError = types.Universe.Lookup("error").Type()
	typeByte  = types.Universe.Lookup("byte").Type()
)

func (g *generator) genFunc(fn *types.Func) {
	g.nonConcrete = false
	sign := fn.Signature()
	if sign.Recv() != nil {
		return
	}
	params := sign.Params()
	results := sign.Results()
	if results == nil || (results.Len() != 1 && results.At(1).Type() != typeError) {
		fmt.Printf("Dropped func %s.%s: must have one return value or a value and an error %v\n", g.cuePkgPath, fn.Name(), sign)
		return
	}

	g.sep()
	fmt.Fprintf(g.w, "{\n")
	defer fmt.Fprintf(g.w, "}")

	fmt.Fprintf(g.w, "Name: %q,\n", fn.Name())

	needCallContext := false
	args := []string{}
	vals := []string{}
	kind := []string{}
	for i := 0; i < params.Len(); i++ {
		param := params.At(i)
		typ := param.Type()
		if typ.String() == "cuelang.org/go/internal/pkg.Schema" {
			needCallContext = true
		}
		methodName := g.callCtxtGetter(typ)
		argKind := g.adtKind(param.Type())
		vals = append(vals, fmt.Sprintf("c.%s(%d)", methodName, len(args)))
		args = append(args, param.Name())
		kind = append(kind, argKind)
	}

	fmt.Fprintf(g.w, "Params: []pkg.Param{\n")
	for _, k := range kind {
		fmt.Fprintf(g.w, "{Kind: %s},\n", k)
	}
	fmt.Fprintf(g.w, "\n},\n")

	fmt.Fprintf(g.w, "Result: %s,\n", g.adtKind(results.At(0).Type()))
	if g.nonConcrete {
		fmt.Fprintf(g.w, "NonConcrete: true,\n")
	}

	argList := strings.Join(args, ", ")
	valList := strings.Join(vals, ", ")
	init := ""
	if len(args) > 0 {
		init = fmt.Sprintf("%s := %s", argList, valList)
	}

	name := fn.Name()
	if needCallContext {
		argList = "c.OpContext(), " + argList

		// Main function is used for Godoc documentation. Once we have proper
		// CUE function signatures, we can remove these stubs.
		// NOTE: this will not work for scripts that are not cased. But this
		// is intended to be a temporary situation anyway.
		name = strings.ToLower(name[:1]) + name[1:]
	}

	fmt.Fprintf(g.w, "Func: func(c *pkg.CallCtxt) {")
	defer fmt.Fprintln(g.w, "},")
	fmt.Fprintln(g.w)
	if init != "" {
		fmt.Fprintln(g.w, init)
	}
	fmt.Fprintln(g.w, "if c.Do() {")
	defer fmt.Fprintln(g.w, "}")
	if results.Len() == 1 {
		fmt.Fprintf(g.w, "c.Ret = %s(%s)", name, argList)
	} else {
		fmt.Fprintf(g.w, "c.Ret, c.Err = %s(%s)", name, argList)
	}
}

// callCtxtGetter returns the name of the [cuelang.org/go/internal/pkg.CallCtxt] method
// which can be used to fetch a parameter of the given type.
func (g *generator) callCtxtGetter(typ types.Type) string {
	switch typ := typ.(type) {
	case *types.Basic:
		return strings.Title(typ.String()) // "int" turns into "Int"
	case *types.Slice:
		switch typ.Elem().String() {
		case "byte":
			return "Bytes"
		case "string":
			return "StringList"
		case "*cuelang.org/go/internal.Decimal":
			return "DecimalList"
		}
		return "List"
	}
	switch typ.String() {
	case "*math/big.Int":
		return "BigInt"
	case "*math/big.Float":
		return "BigFloat"
	case "*cuelang.org/go/internal.Decimal":
		return "Decimal"
	case "cuelang.org/go/internal/pkg.List":
		return "CueList"
	case "cuelang.org/go/internal/pkg.Struct":
		return "Struct"
	case "cuelang.org/go/cue.Value",
		"cuelang.org/go/cue/ast.Expr":
		return "Value"
	case "cuelang.org/go/internal/pkg.Schema":
		g.nonConcrete = true
		return "Schema"
	case "io.Reader":
		return "Reader"
	}
	log.Fatal("callCtxtGetter: unhandled Go type ", typ.String())
	return ""
}

// adtKind provides a Go expression string which describes
// a [cuelang.org/go/internal/core/adt.Kind] value for the given type.
func (g *generator) adtKind(typ types.Type) string {
	// TODO: detect list and structs types for return values.
	switch typ := typ.(type) {
	case *types.Slice:
		if typ.Elem() == typeByte {
			return "adt.BytesKind | adt.StringKind"
		}
		return "adt.ListKind"
	case *types.Map:
		return "adt.StructKind"
	case *types.Basic:
		if typ.Info()&types.IsInteger != 0 {
			return "adt.IntKind"
		}
		if typ.Kind() == types.Float64 {
			return "adt.NumberKind"
		}
		return "adt." + strings.Title(typ.String()) + "Kind" // "bool" turns into "adt.BoolKind"
	}
	switch typ.String() {
	case "error":
		return "adt.BottomKind"
	case "io.Reader":
		return "adt.BytesKind | adt.StringKind"
	case "cuelang.org/go/internal/pkg.Struct":
		return "adt.StructKind"
	case "cuelang.org/go/internal/pkg.List":
		return "adt.ListKind"
	case "*math/big.Int":
		return "adt.IntKind"
	case "*cuelang.org/go/internal.Decimal", "*math/big.Float":
		return "adt.NumberKind"
	case "cuelang.org/go/cue.Value", "cuelang.org/go/cue/ast.Expr", "cuelang.org/go/internal/pkg.Schema":
		return "adt.TopKind" // TODO: can be more precise

	// Some builtin functions return custom types, like [cuelang.org/go/pkg/time.Split].
	// TODO: we can simplify this once the CUE API declarations in ./pkg/...
	// use CUE function signatures to validate their parameters and results.
	case "*cuelang.org/go/pkg/time.Parts":
		return "adt.StructKind"
	}
	log.Fatal("adtKind: unhandled Go type ", typ.String())
	return ""
}

// loadCUEDecls parses a single CUE file from a directory and returns its contents
// as an expression, typically a struct holding all of a file's declarations.
// If there are no CUE files, it returns (nil, nil).
func loadCUEDecls(dir string) (ast.Expr, error) {
	cuefiles, err := filepath.Glob(filepath.Join(dir, "*.cue"))
	if err != nil || len(cuefiles) == 0 {
		return nil, err
	}
	if len(cuefiles) == 0 {
		return nil, nil
	}
	if len(cuefiles) > 1 {
		// Supporting multiple CUE files would require merging declarations.
		return nil, fmt.Errorf("multiple CUE files not supported in this generator")
	}
	src, err := os.ReadFile(cuefiles[0])
	if err != nil {
		return nil, err
	}
	file, err := parser.ParseFile(cuefiles[0], src)
	if err != nil {
		return nil, err
	}
	return internal.ToExpr(file), nil
}