// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build ignore

// AliasGen collects all top-level, exported declarations and
// produces aliases that reference declarations in another package
// that actually implements each declaration.
//
// Usage:
//
//	go run alias_gen.go {TargetPkgPath} {WorkingDir}
//
// Where:
//   - TargetPkgPath is that package path that implements each declaration.
//   - WorkingDir is the directory that contains Go source files.
//     All top-level, exported declarations are collected and transformed
//     into aliases to the equivalent declaration in the target package
//     and written to an alias.go output file.
package main

import (
	"bytes"
	"cmp"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"maps"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
)

func main() {
	targetPkgPath := os.Args[1]
	workingDir := os.Args[2]
	generateAliases(targetPkgPath, workingDir)
}

func generateAliases(targetPkgPath, workingDir string) {
	fset := token.NewFileSet()
	var files []*ast.File
	for _, fi := range mustGet(os.ReadDir(workingDir)) {
		if !strings.HasSuffix(fi.Name(), ".go") || strings.HasSuffix(fi.Name(), "_test.go") || fi.Name() == "alias.go" || fi.Name() == "alias_gen.go" {
			continue
		}
		b := mustGet(os.ReadFile(filepath.Join(workingDir, fi.Name())))
		f := mustGet(parser.ParseFile(fset, fi.Name(), b, parser.ParseComments))
		files = append(files, f)
	}
	slices.SortFunc(files, func(x, y *ast.File) int {
		return cmp.Compare(fset.File(x.Pos()).Name(), fset.File(y.Pos()).Name())
	})

	var aliasFile, aliasDecls bytes.Buffer

	// Print copyright.
	aliasFile.WriteString("// Copyright 2025 The Go Authors. All rights reserved.\n")
	aliasFile.WriteString("// Use of this source code is governed by a BSD-style\n")
	aliasFile.WriteString("// license that can be found in the LICENSE file.\n")
	aliasFile.WriteString("\n")

	// Print build tag.
	aliasFile.WriteString("// Code generated by alias_gen.go; DO NOT EDIT.\n")
	aliasFile.WriteString("\n")
	aliasFile.WriteString("//go:build goexperiment.jsonv2 && go1.25\n")
	aliasFile.WriteString("\n")

	// Print package docs.
	packageName := files[0].Name.String()
	for _, f := range files {
		writeComments(&aliasFile, f.Doc)
	}
	aliasFile.WriteString("package " + packageName + "\n")
	aliasFile.WriteString("\n")

	// Print the imports.
	imports := make(map[string]struct{})
	writeType := func(expr ast.Expr) {
		ast.Walk(astVisitor(func(node ast.Node) bool {
			if sel, ok := node.(*ast.SelectorExpr); ok {
				if id, ok := sel.X.(*ast.Ident); ok {
					switch pkgName := id.String(); pkgName {
					case "io", "bytes":
						imports[pkgName] = struct{}{}
					case "jsontext":
						imports[`encoding/json/jsontext`] = struct{}{}
					default:
						panic(fmt.Sprintf("unknown package %q", id.String()))
					}
				}
			}
			return true
		}), expr)
		mustDo(format.Node(&aliasDecls, fset, expr))
	}
	writeImports := func() {
		aliasFile.WriteString("import (\n")
		imports := append(slices.Collect(maps.Keys(imports)), targetPkgPath)
		slices.Sort(imports)
		for _, pkgPath := range imports {
			aliasFile.WriteString(strconv.Quote(pkgPath) + "\n")
		}
		aliasFile.WriteString(")\n")
		aliasFile.WriteString("\n")
	}

	// Print aliases to every exported top-level declaration.
	for _, f := range files {
		for _, d := range f.Decls {
			switch d := d.(type) {
			case *ast.GenDecl:
				switch d.Tok {
				case token.IMPORT:
				case token.CONST, token.VAR:
					// Check whether there any exported declarations.
					var hasExported bool
					for _, s := range d.Specs {
						for _, name := range s.(*ast.ValueSpec).Names {
							hasExported = hasExported || (name.IsExported() && name.String() != "Internal")
						}
					}
					if !hasExported {
						continue
					}

					// Print the declaration.
					writeComments(&aliasDecls, d.Doc)
					if d.Lparen > 0 {
						aliasDecls.WriteString(d.Tok.String())
						aliasDecls.WriteString(" (\n")
					}
					for _, s := range d.Specs {
						s := s.(*ast.ValueSpec)
						writeComments(&aliasDecls, s.Doc)
						if d.Lparen == 0 {
							aliasDecls.WriteString(d.Tok.String())
							aliasDecls.WriteByte(' ')
						}
						var hasExported bool
						for _, name := range s.Names {
							if name.IsExported() {
								aliasDecls.WriteString(name.String())
								aliasDecls.WriteByte(',')
								hasExported = true
							}
						}
						if !hasExported {
							continue
						}
						trimRight(&aliasDecls, ",")
						aliasDecls.WriteByte('=')
						for _, name := range s.Names {
							if name.IsExported() {
								aliasDecls.WriteString(packageName)
								aliasDecls.WriteByte('.')
								aliasDecls.WriteString(name.String())
								aliasDecls.WriteByte(',')
							}
						}
						trimRight(&aliasDecls, ",")
						aliasDecls.WriteString("\n")
					}
					if d.Rparen > 0 {
						aliasDecls.WriteString(")")
					}
					aliasDecls.WriteString("\n")
				case token.TYPE:
					for _, s := range d.Specs {
						s := s.(*ast.TypeSpec)
						if !s.Name.IsExported() {
							continue
						}
						writeComments(&aliasDecls, d.Doc)
						aliasDecls.WriteString(d.Tok.String())
						aliasDecls.WriteByte(' ')
						aliasDecls.WriteString(s.Name.String())
						aliasDecls.WriteByte('=')
						aliasDecls.WriteString(packageName)
						aliasDecls.WriteByte('.')
						aliasDecls.WriteString(s.Name.String())
						aliasDecls.WriteString("\n")
						aliasDecls.WriteString("\n")
					}
				default:
					panic(fmt.Sprintf("unknown token.Token: %v", d.Tok))
				}
			case *ast.FuncDecl:
				if !d.Name.IsExported() || d.Recv != nil {
					continue // ignore unexported functions or methods
				}

				// Print the comment.
				writeComments(&aliasDecls, d.Doc)
				aliasDecls.WriteString(token.FUNC.String())
				aliasDecls.WriteByte(' ')
				aliasDecls.WriteString(d.Name.String())
				writeFields := func(fields *ast.FieldList, leftDelim, rightDelim byte, withType bool) {
					if fields == nil {
						return
					}
					aliasDecls.WriteByte(leftDelim)
					for i, field := range fields.List {
						for j, name := range field.Names {
							aliasDecls.WriteString(name.String())
							if j < len(field.Names)-1 {
								aliasDecls.WriteByte(',')
							}
						}
						if withType {
							aliasDecls.WriteByte(' ')
							writeType(field.Type)
						} else if _, ok := field.Type.(*ast.Ellipsis); ok {
							aliasDecls.WriteString("...")
						}
						if i < len(fields.List)-1 {
							aliasDecls.WriteByte(',')
						}
					}
					aliasDecls.WriteByte(rightDelim)
				}

				writeFields(d.Type.TypeParams, '[', ']', true)
				writeFields(d.Type.Params, '(', ')', true)
				writeFields(d.Type.Results, '(', ')', true)

				aliasDecls.WriteString("{\n")
				if d.Type.Results != nil {
					aliasDecls.WriteString(token.RETURN.String())
					aliasDecls.WriteByte(' ')
				}
				aliasDecls.WriteString(packageName)
				aliasDecls.WriteByte('.')
				aliasDecls.WriteString(d.Name.String())
				writeFields(d.Type.TypeParams, '[', ']', false)
				writeFields(d.Type.Params, '(', ')', false)
				aliasDecls.WriteString("\n")
				aliasDecls.WriteString("}\n")
				aliasDecls.WriteString("\n")
			default:
				panic(fmt.Sprintf("unknown ast.Decl type: %T", d))
			}
		}
	}
	writeImports()
	aliasFile.Write(aliasDecls.Bytes())

	// Write to the output file.
	b := mustGet(format.Source(aliasFile.Bytes()))
	mustDo(os.WriteFile(filepath.Join(workingDir, "alias.go"), b, 0664))
}

func mustDo(err error) {
	if err != nil {
		panic(err)
	}
}

func mustGet[T any](v T, err error) T {
	mustDo(err)
	return v
}

func writeComments(out *bytes.Buffer, comments *ast.CommentGroup) {
	for line := range strings.Lines(comments.Text()) {
		out.WriteString("// ")
		out.WriteString(line)
	}
}

func trimRight(out *bytes.Buffer, cutset string) {
	out.Truncate(len(bytes.TrimRight(out.Bytes(), cutset)))
}

type astVisitor func(ast.Node) bool

func (f astVisitor) Visit(node ast.Node) ast.Visitor {
	if !f(node) {
		return nil
	}
	return f
}