// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

package main

import (
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/constant"
	"go/parser"
	"go/token"
	"go/types"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
)

// go_asm.h is the header the Go compiler writes for every package that
// carries assembly (the compiler's -asmhdr output): "#define const_NAME
// value" for each package constant, and for each named struct type
// "#define TYPE__size size" plus one "#define TYPE_field offset" per field.
// GOROOT assembly includes it, and a standalone assembler has no compiler
// to have produced it, so gasm generates the equivalent itself: the package
// the .s file lives in is parsed and type-checked here, with the target
// architecture's own sizes, and the same defines are written out.  The
// type-checking GOOS is selected by the caller: a GOOS-specific file
// (sys_darwin_arm64.s) needs its platform's defines, which a header from
// the ambient GOOS silently omits.
//
// The emitter mirrors cmd/compile's dumpasmhdr exactly: constants come out
// as "const_NAME", struct entries as "NAME__size" followed by the fields in
// declaration order, blank names are skipped, and float and complex
// constants are omitted (the assembler carries integers, bools and strings
// only).  Aliases to structs are emitted, generic types are not: they have
// no fixed size.  A define the assembly references but this header does not
// carry surfaces later as the assembler's own "undefined" diagnostic naming
// the define, which is the honest failure.

// goAsmInclude matches the #include "go_asm.h" directive, tolerant of
// whitespace, so the wiring knows which files need a generated header
// before the preprocessor runs and would report the header as missing.
var goAsmInclude = regexp.MustCompile(`(?m)^\s*#\s*include\s+"go_asm\.h"`)

// needsGoAsmHeader reports whether src includes go_asm.h.
func needsGoAsmHeader(src string) bool {
	return goAsmInclude.MatchString(src)
}

// goAsmHeaderResolved reports whether the include of go_asm.h from a file in
// asmDir already resolves: to a header in the package directory itself, or
// in one of the -I directories, the way the preprocessor searches.  Only an
// unresolved include is generated for; a header someone placed by hand is
// the tool the author chose, and it also wins the preprocessor's own search
// order, so generating a second copy would be dead weight at best.
func goAsmHeaderResolved(asmDir string, dirs []string) bool {
	candidates := []string{filepath.Join(asmDir, "go_asm.h")}
	for _, d := range dirs {
		candidates = append(candidates, filepath.Join(d, "go_asm.h"))
	}
	for _, candidate := range candidates {
		if st, err := os.Stat(candidate); err == nil && !st.IsDir() {
			return true
		}
	}
	return false
}

// generateGoAsmHeader type-checks the Go package in pkgDir for goos and
// goarch, writes its go_asm.h equivalent into dir, and returns dir.  An
// empty goos means the ambient one.  The caller owns the directory and its
// removal.
func generateGoAsmHeader(pkgDir, goos, goarch, dir string) (string, error) {
	if goos == "" {
		goos = build.Default.GOOS
	}
	imp := newSourceImporter(goos, goarch)
	if imp.sizes == nil {
		return "", fmt.Errorf("go_asm.h: unknown GOARCH %q", goarch)
	}
	bp, err := imp.ctxt.ImportDir(pkgDir, 0)
	if err != nil {
		return "", fmt.Errorf("go_asm.h for GOARCH %s in %s: %w", goarch, pkgDir, err)
	}
	files, errs := imp.parse(bp)
	if len(errs) > 0 {
		return "", fmt.Errorf("go_asm.h for GOARCH %s in %s: %s", goarch, pkgDir, errorList(errs))
	}
	_, info, errs := imp.checkPackage(bp, files)
	if len(errs) > 0 {
		return "", fmt.Errorf("go_asm.h for GOARCH %s in %s: package does not type-check: %s", goarch, pkgDir, errorList(errs))
	}

	var b strings.Builder
	fmt.Fprintf(&b, "// generated by gasm from package %s (GOOS %s, GOARCH %s)\n\n", bp.Name, goos, goarch)
	// Files in the build's own order and declarations in source order: the
	// same walk the compiler's reader makes, so the header reads the same
	// way the toolchain's does.  Order carries no meaning to the assembler
	// (defines form a table), only to a human diffing against one.
	for _, f := range files {
		for _, decl := range f.Decls {
			gd, ok := decl.(*ast.GenDecl)
			if !ok {
				continue
			}
			for _, spec := range gd.Specs {
				switch gd.Tok {
				case token.CONST:
					vs, ok := spec.(*ast.ValueSpec)
					if !ok {
						continue
					}
					for _, name := range vs.Names {
						emitConst(&b, info.Defs[name], name.Name)
					}
				case token.TYPE:
					ts, ok := spec.(*ast.TypeSpec)
					if !ok {
						continue
					}
					emitStruct(&b, imp.sizes, info.Defs[ts.Name], ts.Name.Name)
				}
			}
		}
	}

	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", fmt.Errorf("go_asm.h for GOARCH %s in %s: %w", goarch, pkgDir, err)
	}
	out := filepath.Join(dir, "go_asm.h")
	if err := os.WriteFile(out, []byte(b.String()), 0o644); err != nil {
		return "", fmt.Errorf("go_asm.h for GOARCH %s in %s: %w", goarch, pkgDir, err)
	}
	return dir, nil
}

// emitConst writes one const define, skipping what the toolchain skips:
// blank names, and float and complex values the assembler has no syntax for.
func emitConst(b *strings.Builder, obj types.Object, name string) {
	c, ok := obj.(*types.Const)
	if !ok || name == "_" {
		return
	}
	switch c.Val().Kind() {
	case constant.Float, constant.Complex, constant.Unknown:
		return
	}
	fmt.Fprintf(b, "#define const_%s %s\n", name, c.Val().ExactString())
}

// emitStruct writes one named struct type's size and field offsets,
// skipping what the toolchain skips: blank names, non-struct types, and
// generic types, whose size depends on their instantiation.
func emitStruct(b *strings.Builder, sizes types.Sizes, obj types.Object, name string) {
	tn, ok := obj.(*types.TypeName)
	if !ok || name == "_" {
		return
	}
	t := types.Unalias(tn.Type())
	// Generic types are spelled *types.Named with a type-parameter list;
	// a plain struct type or an instantiated one carries none.
	if named, ok := t.(*types.Named); ok && named.TypeParams().Len() > 0 {
		return
	}
	st, ok := t.Underlying().(*types.Struct)
	if !ok {
		return
	}
	fmt.Fprintf(b, "#define %s__size %d\n", name, sizes.Sizeof(t))
	fields := make([]*types.Var, st.NumFields())
	for i := range st.NumFields() {
		fields[i] = st.Field(i)
	}
	for i, off := range sizes.Offsetsof(fields) {
		fld := fields[i]
		if fld.Name() == "_" {
			continue
		}
		fmt.Fprintf(b, "#define %s_%s %d\n", name, fld.Name(), off)
	}
}

// errorList renders at most three errors, enough to say what is wrong
// without burying the diagnostic the caller actually reads.
func errorList(errs []error) string {
	if len(errs) > 3 {
		errs = errs[:3]
	}
	msgs := make([]string, len(errs))
	for i, err := range errs {
		msgs[i] = err.Error()
	}
	return strings.Join(msgs, "; ")
}

// sourceImporter type-checks imported packages from source with the target
// architecture's sizes.  go/importer's "source" importer pins the host
// GOARCH, which would lay out imported types (internal/cpu, internal/abi)
// for the wrong target on a cross-architecture header, so the recursion is
// carried here with one build context and one sizes instance per
// architecture.
type sourceImporter struct {
	fset  *token.FileSet
	ctxt  *build.Context
	sizes types.Sizes
	pkgs  map[string]*types.Package
}

// newSourceImporter returns the importer for one target GOOS and GOARCH.
// Cgo is disabled so the file set is deterministic and independent of the
// host's C toolchain: cgo-tagged files drop out of the build exactly as
// they do from a CGO_ENABLED=0 build, whose assembly is what gasm targets.
func newSourceImporter(goos, goarch string) *sourceImporter {
	ctxt := new(build.Context)
	*ctxt = build.Default
	ctxt.GOOS = goos
	ctxt.GOARCH = goarch
	ctxt.CgoEnabled = false
	return &sourceImporter{
		fset:  token.NewFileSet(),
		ctxt:  ctxt,
		sizes: types.SizesFor("gc", goarch),
		pkgs:  map[string]*types.Package{},
	}
}

// Import type-checks one imported package and memoises it.  "unsafe" must
// resolve to go/types' own package, never to the source in GOROOT/src/unsafe:
// the source declares Sizeof and Offsetof as ordinary functions over
// ArbitraryType, and checking against that signature rejects half the
// unsafe arithmetic the gc compiler accepts, which is exactly the divergence
// srcimporter guards against the same way.
func (im *sourceImporter) Import(path string) (*types.Package, error) {
	if path == "unsafe" {
		return types.Unsafe, nil
	}
	if p, ok := im.pkgs[path]; ok {
		return p, nil
	}
	bp, err := im.ctxt.Import(path, "", 0)
	if err != nil {
		return nil, err
	}
	files, errs := im.parse(bp)
	if len(errs) > 0 {
		return nil, errors.New(errorList(errs))
	}
	pkg, _, _ := im.checkPackage(bp, files)
	im.pkgs[path] = pkg
	return pkg, nil
}

// parse reads the build package's Go files.  Import-level failures (no Go
// files for the target, unreadable files) come back as errors, and the
// type-check decides the rest.
func (im *sourceImporter) parse(bp *build.Package) ([]*ast.File, []error) {
	if len(bp.GoFiles) == 0 {
		return nil, []error{fmt.Errorf("no Go source files for GOOS=%s GOARCH=%s", im.ctxt.GOOS, im.ctxt.GOARCH)}
	}
	var (
		files []*ast.File
		errs  []error
	)
	for _, name := range bp.GoFiles {
		f, err := parser.ParseFile(im.fset, filepath.Join(bp.Dir, name), nil, parser.SkipObjectResolution)
		if err != nil {
			errs = append(errs, err)
			continue
		}
		files = append(files, f)
	}
	return files, errs
}

// checkPackage type-checks one package's files with the importer's sizes,
// recording every error: a header from a package that does not type-check
// could silently mis-state an offset, so the caller refuses the header
// rather than trusting it.  The returned Defs map backs the root package's
// emission walk; imports only need the checked package itself.
func (im *sourceImporter) checkPackage(bp *build.Package, files []*ast.File) (*types.Package, *types.Info, []error) {
	var errs []error
	conf := &types.Config{
		Importer: im,
		Sizes:    im.sizes,
		Error:    func(err error) { errs = append(errs, err) },
	}
	info := &types.Info{Defs: map[*ast.Ident]types.Object{}}
	pkg, _ := conf.Check(bp.ImportPath, im.fset, files, info)
	return pkg, info, errs
}

// asmhdrCache generates one go_asm.h per package directory and target
// architecture under one temp root, for callers that assemble many files
// (the corpus audit).  Failures are cached too: a package that does not
// type-check must not be re-checked once per file.
type asmhdrCache struct {
	root string
	dirs map[string]string // "pkgDir\x00goos\x00goarch" -> directory holding go_asm.h
	errs map[string]error
}

func newAsmhdrCache() (*asmhdrCache, error) {
	root, err := os.MkdirTemp("", "gasm-asmhdr")
	if err != nil {
		return nil, err
	}
	return &asmhdrCache{root: root, dirs: map[string]string{}, errs: map[string]error{}}, nil
}

// dirFor returns the directory holding the generated go_asm.h for pkgDir
// under goos and goarch, generating it on first use.  An empty goos means
// the ambient one, resolved here so that one package cannot generate twice
// under an explicit and an implicit spelling of the same GOOS.
func (c *asmhdrCache) dirFor(pkgDir, goos, goarch string) (string, error) {
	if goos == "" {
		goos = build.Default.GOOS
	}
	key := pkgDir + "\x00" + goos + "\x00" + goarch
	if dir, ok := c.dirs[key]; ok {
		return dir, nil
	}
	if err, ok := c.errs[key]; ok {
		return "", err
	}
	dir := filepath.Join(c.root, fmt.Sprintf("h%d_%s_%s", len(c.dirs), goos, goarch))
	if _, err := generateGoAsmHeader(pkgDir, goos, goarch, dir); err != nil {
		c.errs[key] = err
		return "", err
	}
	c.dirs[key] = dir
	return dir, nil
}

// close removes the temp root.
func (c *asmhdrCache) close() { os.RemoveAll(c.root) }

// ensureGoAsmHeader prepares the include directory a file that includes
// go_asm.h needs: the generated header for the package in path's directory,
// for the file's target GOOS and architecture.  It reports a usage error
// when the architecture cannot be determined, and passes through the
// generator's diagnostics, which name the package.
func ensureGoAsmHeader(path string, target arch.Arch, goos string, cache *asmhdrCache) (string, func(), error) {
	if path == "-" {
		return "", nil, errors.New("cannot generate go_asm.h for standard input (no package directory)")
	}
	if target == arch.Unknown {
		return "", nil, errors.New("a file that includes go_asm.h needs a target architecture: name the file _<arch>.s or pass -GOARCH")
	}
	if cache != nil {
		dir, err := cache.dirFor(filepath.Dir(path), goos, goarchName(target))
		return dir, func() {}, err
	}
	root, err := os.MkdirTemp("", "gasm-asmhdr")
	if err != nil {
		return "", nil, err
	}
	dir, err := generateGoAsmHeader(filepath.Dir(path), goos, goarchName(target), root)
	if err != nil {
		os.RemoveAll(root)
		return "", nil, err
	}
	return dir, func() { os.RemoveAll(root) }, nil
}