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

package verify

import (
	"fmt"
	"os"

	"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)
}

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