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

package verify

import (
	"fmt"
	"math/rand"
	"os"
	"runtime"

	"sourcedock.dev/petrbalvin/gasm-devkit/asm"
	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
)

// Kernel is a JIT-loaded assembly image ready for direct invocation.
// It wraps an executable memory mapping and the function layout metadata
// needed to marshal ABI0 calls.
type Kernel struct {
	exec  *Executable
	img   *asm.Image
	funcs map[string]int // function name → index into img.Funcs
}

// Load parses, assembles and maps a .s file into executable memory.
// The returned Kernel is ready for Call.  The caller must call Close to
// release the mapping.
func Load(path string) (*Kernel, error) {
	src, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("verify: %w", err)
	}
	return LoadSource(path, string(src))
}

// LoadSource parses, assembles and maps assembly source into executable memory.
func LoadSource(filename, src string) (*Kernel, error) {
	file, errs := parser.Parse(filename, src)
	if len(errs) > 0 {
		return nil, fmt.Errorf("verify: parse %s: %v", filename, errs[0])
	}
	return LoadAST(file)
}

// LoadAST assembles a parsed AST file and maps the result into executable
// memory.
func LoadAST(file *ast.File) (*Kernel, error) {
	img, err := asm.AssembleFile(file)
	if err != nil {
		return nil, fmt.Errorf("verify: assemble: %w", err)
	}
	if len(img.Externals) > 0 {
		return nil, fmt.Errorf("verify: unresolved external symbols: %v", img.Externals)
	}
	code := img.Bytes()
	exec, err := Map(code)
	if err != nil {
		return nil, err
	}
	funcs := make(map[string]int, len(img.Funcs))
	for i, f := range img.Funcs {
		funcs[f.Name] = i
	}
	return &Kernel{exec: exec, img: img, funcs: funcs}, nil
}

// Func returns the layout metadata for the named function.
func (k *Kernel) Func(name string) (asm.FuncLayout, error) {
	idx, ok := k.funcs[name]
	if !ok {
		return asm.FuncLayout{}, fmt.Errorf("verify: function %q not found", name)
	}
	return k.img.Funcs[idx], nil
}

// FuncNames returns the names of all functions in the kernel, in source order.
func (k *Kernel) FuncNames() []string {
	names := make([]string, len(k.img.Funcs))
	for i, f := range k.img.Funcs {
		names[i] = f.Name
	}
	return names
}

// CallFunc invokes the named function with the given ABI0 argument block.
// The arg block is the raw bytes of the function's argument/result area
// (as declared by the TEXT $frame-args suffix).  Returns the arg block
// after the call (with any results written back by the function).
func (k *Kernel) CallFunc(name string, args []byte) ([]byte, error) {
	idx, ok := k.funcs[name]
	if !ok {
		return nil, fmt.Errorf("verify: function %q not found", name)
	}
	fl := k.img.Funcs[idx]
	if len(args) < fl.Args {
		return nil, fmt.Errorf("verify: %s: arg block too small: got %d, need %d", name, len(args), fl.Args)
	}
	fnAddr := k.exec.FuncAddr(fl.Offset)
	return Call(fnAddr, args)
}

// CallFuncChecked invokes the named function with ABI sentinels and a
// red-zone canary, returning the argument block and an ABIReport that
// records any callee-saved register or red-zone violations.
func (k *Kernel) CallFuncChecked(name string, args []byte) ([]byte, ABIReport, error) {
	idx, ok := k.funcs[name]
	if !ok {
		return nil, ABIReport{}, fmt.Errorf("verify: function %q not found", name)
	}
	fl := k.img.Funcs[idx]
	if len(args) < fl.Args {
		return nil, ABIReport{}, fmt.Errorf("verify: %s: arg block too small: got %d, need %d", name, len(args), fl.Args)
	}
	fnAddr := k.exec.FuncAddr(fl.Offset)
	return CallChecked(fnAddr, args)
}

// FuzzFuncCheckedByName is like FuzzFuncChecked but extracts the signature
// from the source code internally.
func (k *Kernel) FuzzFuncCheckedByName(name, src string, iterations int, seed int64) FuzzResult {
	result := FuzzResult{Func: name, Iterations: iterations}
	sig, ok := ExtractSignatures(src)[name]
	if !ok {
		result.Mismatches = iterations
		result.FirstFail = "no // func signature found"
		return result
	}
	return k.FuzzFuncChecked(name, sig, iterations, seed)
}

// FuzzFuncChecked combines fuzzing with ABI checks: it generates varied
// inputs and verifies that callee-saved registers and the red zone are
// preserved even on deep execution paths (not just early exits).
func (k *Kernel) FuzzFuncChecked(name string, sig funcSig, iterations int, seed int64) FuzzResult {
	result := FuzzResult{Func: name, Iterations: iterations}
	rng := rand.New(rand.NewSource(seed))

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

	violations := 0
	for i := range iterations {
		gasmArgs, _, bufs, _ := genDualArgs(rng, sig, fl.Args)
		result.CrashInput = gasmArgs

		_, report, err := k.CallFuncChecked(name, gasmArgs)
		if err != nil {
			result.Mismatches++
			if result.FirstFail == "" {
				result.FirstFail = fmt.Sprintf("iter %d: call: %v", i, err)
			}
			runtime.KeepAlive(bufs)
			continue
		}

		if !report.OK() {
			violations++
			result.Mismatches++
			if result.FirstFail == "" {
				result.FirstFail = fmt.Sprintf("iter %d: %s", i, report.String())
			}
		} else {
			result.Matches++
		}
		runtime.KeepAlive(bufs)
	}

	if violations > 0 && result.FirstFail == "" {
		result.FirstFail = fmt.Sprintf("%d ABI violations across %d iterations", violations, iterations)
	}
	return result
}

// Close releases the executable mapping.
func (k *Kernel) Close() {
	if k.exec != nil {
		k.exec.Unmap()
	}
}

// Image returns the assembled image (code + data + metadata).
func (k *Kernel) Image() *asm.Image {
	return k.img
}