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

package lint

import (
	"go/ast"
	"go/parser"
	"go/token"
	"strconv"
	"strings"
)

// abiExpectedArgSize computes the argument-area size (parameters plus results,
// laid out with Go's alignment rules on a 64-bit target) from the `// func …`
// signature in a TEXT function's doc comment.  It returns ok=false when there
// is no parseable signature or it uses a type whose size cannot be determined
// (a named type), so the caller can skip the check rather than guess.
//
// The signature is parsed with the standard library's Go parser, so every
// legal signature form (shared names such as `left, right []int32`, nested
// pointers, arrays, structs) is handled correctly.
func abiExpectedArgSize(doc string) (int64, bool) {
	sig := signatureLine(doc)
	if sig == "" {
		return 0, false
	}
	fn, ok := parseSignature(sig)
	if !ok || fn == nil {
		return 0, false
	}
	return signatureSize(fn.Type.Params, fn.Type.Results)
}

// abiParamNames returns the parameter names declared by the `// func …`
// signature in a doc comment, in declaration order.  Shared names
// (`left, right []int32`) expand to one entry per name; a nameless parameter
// (`[]int32`) contributes an empty placeholder so offsets stay aligned.
func abiParamNames(doc string) ([]string, bool) {
	sig := signatureLine(doc)
	if sig == "" {
		return nil, false
	}
	fn, ok := parseSignature(sig)
	if !ok || fn == nil || fn.Type.Params == nil {
		return nil, false
	}
	var names []string
	for _, field := range fn.Type.Params.List {
		if len(field.Names) == 0 {
			names = append(names, "")
			continue
		}
		for _, n := range field.Names {
			names = append(names, n.Name)
		}
	}
	return names, true
}

// parseSignature parses a `func …` line into a Go FuncDecl.
func parseSignature(sig string) (*ast.FuncDecl, bool) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "sig.go", "package p\n"+sig+" {}\n", 0)
	if err != nil || len(f.Decls) == 0 {
		return nil, false
	}
	fn, ok := f.Decls[0].(*ast.FuncDecl)
	if !ok || fn.Type == nil {
		return nil, false
	}
	return fn, true
}

// signatureLine returns the first `func …` line from a doc comment, trimmed.
func signatureLine(doc string) string {
	for _, line := range strings.Split(doc, "\n") {
		if t := strings.TrimSpace(line); strings.HasPrefix(t, "func ") {
			return t
		}
	}
	return ""
}

// signatureSize lays out the parameters and results and returns the total byte
// size of the argument area, matching Go's ABI0 stack layout: parameters are
// laid out first, then the result area begins on a word (8-byte) boundary.
func signatureSize(params, results *ast.FieldList) (int64, bool) {
	paramsSize, _, ok := fieldsSizeAlign(params)
	if !ok {
		return 0, false
	}
	resultsSize, _, ok := fieldsSizeAlign(results)
	if !ok {
		return 0, false
	}
	// With no results the argument area is exactly the parameter size.  When
	// there are results, the result area begins on a word (8-byte) boundary
	// after the parameters (Go's ABI0 stack layout).
	if resultsSize == 0 {
		return int64(paramsSize), true
	}
	return int64(alignUp(paramsSize, 8) + resultsSize), true
}

// fieldsSizeAlign lays out a field list sequentially (each field aligned to its
// own alignment) and returns the total size and the maximum field alignment.
func fieldsSizeAlign(list *ast.FieldList) (size, align int, ok bool) {
	if list == nil {
		return 0, 1, true
	}
	offset, maxAlign := 0, 1
	for _, field := range list.List {
		es, ea, fieldOK := typeSizeAlign(field.Type)
		if !fieldOK {
			return 0, 0, false
		}
		n := len(field.Names)
		if n == 0 {
			n = 1
		}
		for range n {
			offset = alignUp(offset, ea)
			offset += es
		}
		if ea > maxAlign {
			maxAlign = ea
		}
	}
	return offset, maxAlign, true
}

// basicSizes maps built-in type names to {size, align} on a 64-bit target.
var basicSizes = map[string][2]int{
	"bool": {1, 1}, "byte": {1, 1}, "int8": {1, 1}, "uint8": {1, 1},
	"int16": {2, 2}, "uint16": {2, 2},
	"int32": {4, 4}, "uint32": {4, 4}, "float32": {4, 4},
	"int": {8, 8}, "int64": {8, 8}, "uint": {8, 8}, "uint64": {8, 8},
	"uintptr": {8, 8}, "float64": {8, 8},
	"complex64": {8, 4}, "complex128": {16, 8},
	"string": {16, 8}, "any": {16, 8}, "error": {16, 8},
}

// typeSizeAlign returns the size and alignment in bytes of a type expression,
// or ok=false when the size cannot be determined (an unknown named type).
func typeSizeAlign(e ast.Expr) (size, align int, ok bool) {
	switch t := e.(type) {
	case *ast.Ident:
		if sa, found := basicSizes[t.Name]; found {
			return sa[0], sa[1], true
		}
		return 0, 0, false // named type of unknown size
	case *ast.SelectorExpr:
		if pkg, isIdent := t.X.(*ast.Ident); isIdent && pkg.Name == "unsafe" && t.Sel.Name == "Pointer" {
			return 8, 8, true
		}
		return 0, 0, false
	case *ast.ParenExpr:
		return typeSizeAlign(t.X)
	case *ast.StarExpr:
		return 8, 8, true // pointer
	case *ast.MapType, *ast.ChanType, *ast.FuncType:
		return 8, 8, true // map / chan / func are pointer-sized
	case *ast.InterfaceType:
		return 16, 8, true
	case *ast.Ellipsis:
		return 24, 8, true // variadic parameter is a slice
	case *ast.ArrayType:
		if t.Len == nil {
			return 24, 8, true // slice header
		}
		n, lenOK := arrayLength(t.Len)
		es, ea, elemOK := typeSizeAlign(t.Elt)
		if !lenOK || !elemOK {
			return 0, 0, false
		}
		return n * es, ea, true
	case *ast.StructType:
		return structSizeAlign(t.Fields)
	}
	return 0, 0, false
}

// structSizeAlign lays out a struct's fields and returns its size (rounded up
// to its alignment) and alignment.
func structSizeAlign(fields *ast.FieldList) (size, align int, ok bool) {
	size, align, ok = fieldsSizeAlign(fields)
	if !ok {
		return 0, 0, false
	}
	return alignUp(size, align), align, true
}

// arrayLength evaluates a constant array-length expression (a literal, for the
// kernels this toolkit targets).
func arrayLength(e ast.Expr) (int, bool) {
	if lit, ok := e.(*ast.BasicLit); ok && lit.Kind == token.INT {
		if v, err := strconv.Atoi(lit.Value); err == nil {
			return v, true
		}
	}
	return 0, false
}

func alignUp(offset, align int) int {
	if align <= 1 {
		return offset
	}
	return (offset + align - 1) &^ (align - 1)
}