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

package verify

import (
	"fmt"
	"math/rand"
	"regexp"
	"strconv"
	"strings"
	"unsafe"
)

// FuzzResult reports the outcome of a differential fuzz campaign for one
// function.
type FuzzResult struct {
	Func       string
	Iterations int
	Matches    int
	Mismatches int
	FirstFail  string // description of the first mismatch ("" if none)
	CrashInput []byte // input that caused the last crash/mismatch (nil if none)
}

// OK returns true when all iterations matched.
func (r FuzzResult) OK() bool { return r.Mismatches == 0 }

// String returns a human-readable summary.
func (r FuzzResult) String() string {
	if r.OK() {
		return fmt.Sprintf("%s: %d/%d iterations match", r.Func, r.Matches, r.Iterations)
	}
	s := fmt.Sprintf("%s: %d/%d match, %d MISMATCH — %s",
		r.Func, r.Matches, r.Iterations, r.Mismatches, r.FirstFail)
	if len(r.CrashInput) > 0 {
		s += fmt.Sprintf("\n  input: %x", r.CrashInput)
	}
	return s
}

// funcSig is a parsed // func signature from the assembly source.
type funcSig struct {
	name    string
	params  []param
	results []param
}

type param struct {
	name string
	typ  string // "[]byte", "[]int32", "int", "*[32]uint16", etc.
}

// funcSigRe matches the conventional "// func name(...)" comment.
var funcSigRe = regexp.MustCompile(`^//\s*func\s+(\w+)\(([^)]*)\)\s*(.*)$`)

// parseFuncSig extracts the function signature from a "// func ..." comment.
func parseFuncSig(comment string) (funcSig, bool) {
	m := funcSigRe.FindStringSubmatch(strings.TrimSpace(comment))
	if m == nil {
		return funcSig{}, false
	}
	sig := funcSig{name: m[1]}
	sig.params = parseParams(m[2])
	// Results may be "(a int, b int)" or "int" or "(int, error)".
	res := strings.TrimSpace(m[3])
	res = strings.TrimPrefix(res, "(")
	res = strings.TrimSuffix(res, ")")
	if res != "" {
		sig.results = parseParams(res)
	}
	return sig, true
}

// parseParams splits "a []byte, b []int32" into typed parameters, handling
// shared types ("a, b []int32").
func parseParams(s string) []param {
	s = strings.TrimSpace(s)
	if s == "" {
		return nil
	}
	var out []param
	for _, field := range strings.Split(s, ",") {
		field = strings.TrimSpace(field)
		if field == "" {
			continue
		}
		parts := strings.Fields(field)
		if len(parts) == 1 {
			// Unnamed: "int" or "[]byte".
			out = append(out, param{typ: parts[0]})
		} else {
			// Named: "a []byte" or shared "a, b []int32" (handled by the
			// comma split above — "a" alone means the type follows in the
			// next field; this is a simplification that covers the common
			// case where each param has its own type).
			out = append(out, param{name: parts[0], typ: parts[1]})
		}
	}
	return out
}

// ExtractSignatures scans assembly source for "// func name(...)" comments
// that immediately precede a TEXT directive, and returns the parsed
// signatures keyed by the function's short name.
func ExtractSignatures(src string) map[string]funcSig {
	lines := strings.Split(src, "\n")
	sigs := make(map[string]funcSig)
	var comments []string
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "//") {
			comments = append(comments, trimmed)
			continue
		}
		if strings.HasPrefix(trimmed, "TEXT") {
			// Search the comment block for the // func line.
			for _, c := range comments {
				if sig, ok := parseFuncSig(c); ok {
					sigs[sig.name] = sig
					break
				}
			}
			comments = nil
			continue
		}
		if trimmed != "" {
			comments = nil
		}
	}
	return sigs
}

// FuzzFunc runs a differential fuzz campaign: it JIT-executes both the
// gasm-assembled and the go-tool-asm-assembled versions of the named
// function with random inputs derived from the // func signature, and
// compares the output argument area bit-for-bit.
//
// The signature comment must appear immediately above the TEXT directive
// in the source (the conventional Go assembly layout).
func (k *Kernel) FuzzFunc(name string, sig funcSig, goCode []byte, iterations int, seed int64) FuzzResult {
	result := FuzzResult{Func: name, Iterations: iterations}

	rng := rand.New(rand.NewSource(seed))

	// Map the Go-assembled code into a second executable region.
	goExec, err := Map(goCode)
	if err != nil {
		result.Mismatches = iterations
		result.FirstFail = fmt.Sprintf("map go code: %v", err)
		return result
	}
	defer goExec.Unmap()

	fl, err := k.Func(name)
	if err != nil {
		result.Mismatches = iterations
		result.FirstFail = err.Error()
		return result
	}

	for i := 0; i < iterations; i++ {
		// Generate inputs and build TWO independent arg blocks (one per
		// version) so that functions which write to their arguments
		// (e.g. histogram increments) don't corrupt the other's input.
		gasmArgs, goArgs, bufs := genDualArgs(rng, sig, fl.Args)

		// Save the current input for crash diagnostics.
		result.CrashInput = gasmArgs

		// Call the gasm version.
		gasmOut, err := k.CallFunc(name, gasmArgs)
		if err != nil {
			result.Mismatches++
			if result.FirstFail == "" {
				result.FirstFail = fmt.Sprintf("iter %d: gasm call: %v", i, err)
			}
			releaseBufs(bufs)
			continue
		}

		// Call the Go version (same function, independent buffers).
		goOut, err := Call(goExec.FuncAddr(0), goArgs)
		if err != nil {
			result.Mismatches++
			if result.FirstFail == "" {
				result.FirstFail = fmt.Sprintf("iter %d: go call: %v", i, err)
			}
			releaseBufs(bufs)
			continue
		}

		// Compare only the result area (after all input parameters).
		// Pointers in the arg block differ (separate buffers), so we
		// compare from resultOff to the end.
		resultOff := paramsSize(sig)
		gasmRes := gasmOut[resultOff:]
		goRes := goOut[resultOff:]
		if !equalBytes(gasmRes, goRes) {
			result.Mismatches++
			if result.FirstFail == "" {
				result.FirstFail = fmt.Sprintf("iter %d: output mismatch at result offset %d", i, resultOff)
			}
		} else {
			result.Matches++
		}
		releaseBufs(bufs)
	}
	return result
}

// genDualArgs generates two independent ABI0 argument blocks (for gasm and
// go) with identical logical content but separate backing buffers, so that
// functions which write to their arguments don't corrupt the other's input.
func genDualArgs(rng *rand.Rand, sig funcSig, argSize int) (gasmArgs, goArgs []byte, bufs [][]byte) {
	gasmArgs = make([]byte, argSize)
	goArgs = make([]byte, argSize)
	off := 0
	sliceIdx := 0

	for _, p := range sig.params {
		switch {
		case strings.HasPrefix(p.typ, "[]"):
			elemSize := elemSizeFor(p.typ)
			n := 1 + rng.Intn(127)
			var declaredLen int
			if sliceIdx == 0 {
				declaredLen = n
			} else {
				declaredLen = n + 512
			}
			// Allocate a buffer comfortably larger than declaredLen*elemSize so
			// that SIMD over-reads and functions that write slightly past len
			// (e.g. decoders that trust len(src)) never touch unmapped memory.
			bufBytes := declaredLen*elemSize + 8192
			// Two independent buffers with identical random content.
			buf1 := make([]byte, bufBytes)
			buf2 := make([]byte, bufBytes)
			rng.Read(buf1[:n*elemSize])
			copy(buf2, buf1)
			bufs = append(bufs, buf1, buf2)
			putPtr(gasmArgs, off, unsafe.Pointer(&buf1[0]))
			putPtr(goArgs, off, unsafe.Pointer(&buf2[0]))
			// len and cap both equal declaredLen — the buffer is guaranteed
			// to hold at least declaredLen elements plus safety margin.
			putU64(gasmArgs, off+8, uint64(declaredLen))
			putU64(gasmArgs, off+16, uint64(declaredLen))
			putU64(goArgs, off+8, uint64(declaredLen))
			putU64(goArgs, off+16, uint64(declaredLen))
			off += 24
			sliceIdx++

		case strings.HasPrefix(p.typ, "*["):
			nElem := arrayLen(p.typ)
			elem := elemSizeFor("[]" + p.typ[strings.Index(p.typ, "]")+1:])
			size := nElem * elem
			if size < 8 {
				size = 8
			}
			buf1 := make([]byte, size)
			buf2 := make([]byte, size)
			rng.Read(buf1)
			copy(buf2, buf1)
			bufs = append(bufs, buf1, buf2)
			putPtr(gasmArgs, off, unsafe.Pointer(&buf1[0]))
			putPtr(goArgs, off, unsafe.Pointer(&buf2[0]))
			off += 8

		case p.typ == "int" || p.typ == "uint" || p.typ == "int64" || p.typ == "uint64":
			v := uint64(rng.Intn(256))
			putU64(gasmArgs, off, v)
			putU64(goArgs, off, v)
			off += 8

		default:
			v := rng.Uint64()
			putU64(gasmArgs, off, v)
			putU64(goArgs, off, v)
			off += 8
		}
	}
	return gasmArgs, goArgs, bufs
}

func elemSizeFor(sliceType string) int {
	switch strings.TrimPrefix(sliceType, "[]") {
	case "byte", "uint8", "int8":
		return 1
	case "uint16", "int16":
		return 2
	case "uint32", "int32", "float32":
		return 4
	case "uint64", "int64", "float64":
		return 8
	default:
		return 8
	}
}

// paramsSize returns the ABI0 stack size occupied by the input parameters.
func paramsSize(sig funcSig) int {
	size := 0
	for _, p := range sig.params {
		switch {
		case strings.HasPrefix(p.typ, "[]"):
			size += 24 // slice header
		case strings.HasPrefix(p.typ, "*["):
			size += 8 // pointer
		case p.typ == "bool":
			size += 1
		default:
			size += 8 // int, uint, etc.
		}
	}
	return size
}

func arrayLen(typ string) int {
	// "*[32]uint16" → 32
	start := strings.Index(typ, "[")
	end := strings.Index(typ, "]")
	if start < 0 || end < 0 || end <= start {
		return 1
	}
	n, _ := strconv.Atoi(typ[start+1 : end])
	if n <= 0 {
		n = 1
	}
	return n
}

func putPtr(buf []byte, off int, p unsafe.Pointer) {
	if off+8 <= len(buf) {
		u64 := uint64(uintptr(p))
		buf[off] = byte(u64)
		buf[off+1] = byte(u64 >> 8)
		buf[off+2] = byte(u64 >> 16)
		buf[off+3] = byte(u64 >> 24)
		buf[off+4] = byte(u64 >> 32)
		buf[off+5] = byte(u64 >> 40)
		buf[off+6] = byte(u64 >> 48)
		buf[off+7] = byte(u64 >> 56)
	}
}

func putU64(buf []byte, off int, v uint64) {
	if off+8 <= len(buf) {
		buf[off] = byte(v)
		buf[off+1] = byte(v >> 8)
		buf[off+2] = byte(v >> 16)
		buf[off+3] = byte(v >> 24)
		buf[off+4] = byte(v >> 32)
		buf[off+5] = byte(v >> 40)
		buf[off+6] = byte(v >> 48)
		buf[off+7] = byte(v >> 56)
	}
}

func equalBytes(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func releaseBufs(bufs [][]byte) {
	// Keep buffers alive until after the call; nothing to free in Go,
	// but this prevents the compiler from collecting them too early.
	_ = bufs
}