core/plugin/wasmguest/dispatch.go

package wasmguest

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sort"
	"time"

	abiv1 "git.dev.alexdunmow.com/block/core/abi/v1"
	"git.dev.alexdunmow.com/block/core/blocks"
	"git.dev.alexdunmow.com/block/core/plugin"
	"git.dev.alexdunmow.com/block/core/plugin/wasmguest/bnwasm"
	"git.dev.alexdunmow.com/block/core/plugin/wasmguest/caps"
	"github.com/google/uuid"
	"google.golang.org/protobuf/proto"
)

// capTransport is the guest→host capability transport handed to the
// capability stubs (package caps). hostcalls.go (wasip1) binds it to CallHost in an init;
// on native builds it stays nil, so capability calls report "no host
// transport" rather than reaching a boundary that isn't there. Tests exercise
// the stubs directly with their own fake transport.
var capTransport caps.CallFunc

// current is the guest runtime installed by Serve. The bn_invoke export
// (exports.go, wasip1) dispatches through it.
var current *guest

// Serve installs a plugin registration as the wasm guest runtime. It is
// non-blocking: plugins call it from an init func (see doc.go) so it runs
// during the module's `_initialize`; all work then arrives via bn_invoke.
func Serve(reg plugin.PluginRegistration) {
	current = newGuest(reg)
}

// HostServices returns the CoreServices bound to THIS guest instance. Unlike
// the deps handed to Load/HTTPHandler/JobHandlers — which the host delivers on
// only one pooled instance (Load) or lazily on first use — these services are
// bound at `_initialize` on EVERY pooled instance (newGuest). The interface
// fields are host_call-backed capability stubs and Pool is a live db.* bridge.
//
// This is the escape hatch for entry points the ABI does not hand deps to —
// most importantly block/template render funcs (HOOK_RENDER_BLOCK gives the
// closure only ctx+content, no services). A DB-backed block reads its pool from
// HostServices().Pool so it works on any pooled instance, not just the one that
// happened to run Load. In a native build (Serve never called — no wasip1
// _initialize) current is nil and this returns the zero CoreServices, so guest
// code must nil-check Pool (native tests inject a real pool by other means).
func HostServices() plugin.CoreServices {
	if current == nil {
		return plugin.CoreServices{}
	}
	return current.services
}

// guest adapts a plugin.PluginRegistration to ABI hook calls.
type guest struct {
	reg         plugin.PluginRegistration
	blocks      *captureBlockRegistry
	templates   *captureTemplateRegistry
	registerErr error

	// services is the CoreServices value handed to the registration's
	// function fields (HTTPHandler, JobHandlers, Load). Its interface fields
	// are caps host_call-backed stubs (package caps); LOAD fills
	// AppURL/MediaPath from HostConfig.
	services plugin.CoreServices

	// rag is the guest-side RAG stub inside services. It records
	// RegisterContentFetcher calls (dispatched via HOOK_RAG_FETCH) while its
	// Query/OnContentChanged marshal out as capability calls.
	rag *caps.RAGStub

	jobHandlers     map[string]plugin.JobHandlerFunc
	jobHandlersInit bool
	httpHandler     http.Handler
	httpHandlerInit bool
}

// newGuest captures the registration's Register pass so block/template
// functions are addressable before any hook fires.
func newGuest(reg plugin.PluginRegistration) *guest {
	g := &guest{
		reg:       reg,
		blocks:    newCaptureBlockRegistry(),
		templates: newCaptureTemplateRegistry(),
	}
	g.services = caps.NewCoreServices(capTransport)
	// Pool is left zero by NewCoreServices (it does not cross as a capability
	// call); bind the db.* driver over the same transport so deps.Pool keeps
	// working for plugin sqlc code. A nil transport (native/DESCRIBE) yields a
	// Pool whose calls fail cleanly, matching the capability stubs.
	g.services.Pool = bnwasm.NewPool(bnwasm.Transport(capTransport))
	g.rag, _ = g.services.RAGService.(*caps.RAGStub)
	g.registerErr = runRegister(reg, g.templates, g.blocks)
	return g
}

// runRegister invokes Register (or RegisterWithProvisioner with a no-op
// provisioner — provisioning itself runs host-side at load) against capture
// registries, converting panics to errors.
func runRegister(reg plugin.PluginRegistration, tr *captureTemplateRegistry, br *captureBlockRegistry) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("panic in Register: %v", r)
		}
	}()
	// Give this registration a clean tag/filter slate. blocks.RegisterTag /
	// RegisterFilter write a package-level registry (one plugin owns the wasm
	// module), so resetting here makes newGuest deterministic: after Register,
	// the registry holds exactly this plugin's tags/filters for DESCRIBE and
	// HOOK_RENDER_TAG / HOOK_APPLY_FILTER dispatch.
	blocks.ResetRegisteredTagsAndFilters()
	// Match the .so loader: block keys are plugin-prefixed and override
	// sources stamped via PluginBlockRegistry.
	pbr := plugin.NewPluginBlockRegistry(br, reg.Name)
	switch {
	case reg.Register != nil:
		return reg.Register(tr, pbr)
	case reg.RegisterWithProvisioner != nil:
		return reg.RegisterWithProvisioner(tr, pbr, noopProvisioner{})
	default:
		return nil
	}
}

// invoke handles one bn_invoke call: (hook id, serialized InvokeRequest) →
// serialized InvokeResponse. It never panics — plugin-level panics come back
// as ABI_ERROR_CODE_INTERNAL so the instance stays usable; traps are
// reserved for runtime corruption.
func (g *guest) invoke(hookID uint32, req []byte) (out []byte) {
	defer func() {
		if r := recover(); r != nil {
			out = errorResponse(abiv1.AbiErrorCode_ABI_ERROR_CODE_INTERNAL, fmt.Sprintf("panic: %v", r))
		}
	}()

	env := &abiv1.InvokeRequest{}
	if err := proto.Unmarshal(req, env); err != nil {
		return errorResponse(abiv1.AbiErrorCode_ABI_ERROR_CODE_DECODE, "invalid InvokeRequest: "+err.Error())
	}
	if uint32(env.GetHook()) != hookID {
		return errorResponse(abiv1.AbiErrorCode_ABI_ERROR_CODE_DECODE,
			fmt.Sprintf("hook id mismatch: export arg %d, envelope %d", hookID, env.GetHook()))
	}
	if g.registerErr != nil {
		return errorResponse(abiv1.AbiErrorCode_ABI_ERROR_CODE_INTERNAL, "plugin Register failed: "+g.registerErr.Error())
	}

	ctx := context.Background()
	if d := env.GetDeadlineMs(); d > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Duration(d)*time.Millisecond)
		defer cancel()
	}

	var (
		payload proto.Message
		abiErr  *abiv1.AbiError
	)
	switch env.GetHook() {
	case abiv1.Hook_HOOK_DESCRIBE:
		payload, abiErr = g.describe(env.GetPayload())
	case abiv1.Hook_HOOK_RENDER_BLOCK:
		payload, abiErr = g.renderBlock(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_RENDER_TEMPLATE:
		payload, abiErr = g.renderTemplate(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_RENDER_TAG:
		payload, abiErr = g.renderTag(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_APPLY_FILTER:
		payload, abiErr = g.applyFilter(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_HANDLE_HTTP:
		payload, abiErr = g.handleHTTP(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_JOB:
		payload, abiErr = g.runJob(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_LOAD:
		payload, abiErr = g.load(env.GetPayload())
	case abiv1.Hook_HOOK_UNLOAD:
		payload, abiErr = g.unload(ctx)
	case abiv1.Hook_HOOK_RAG_FETCH:
		payload, abiErr = g.ragFetch(ctx, env.GetPayload())
	case abiv1.Hook_HOOK_MEDIA_HOOK:
		payload, abiErr = g.mediaHook(ctx, env.GetPayload())
	default:
		abiErr = &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: fmt.Sprintf("unknown hook %d", env.GetHook()),
		}
	}

	resp := &abiv1.InvokeResponse{Error: abiErr}
	if abiErr == nil && payload != nil {
		b, err := proto.Marshal(payload)
		if err != nil {
			return errorResponse(abiv1.AbiErrorCode_ABI_ERROR_CODE_INTERNAL, "marshal response payload: "+err.Error())
		}
		resp.Payload = b
	}
	b, err := proto.Marshal(resp)
	if err != nil {
		return errorResponse(abiv1.AbiErrorCode_ABI_ERROR_CODE_INTERNAL, "marshal InvokeResponse: "+err.Error())
	}
	return b
}

// --- hook handlers ---

func (g *guest) renderBlock(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.RenderBlockRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("RenderBlockRequest", err)
	}
	fn, ok := g.blocks.lookup(req.GetBlockKey(), req.GetRenderContext().GetTemplateKey())
	if !ok {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "unknown block key " + req.GetBlockKey(),
		}
	}
	ctx = contextFromRenderContext(ctx, req.GetRenderContext())
	content := unmarshalMap(req.GetContentJson())
	if content == nil {
		content = map[string]any{}
	}
	out := fn(ctx, content)
	// A block may return EITHER final HTML or a "powered" result
	// (blocks.PoweredBlock): a template + data the HOST renders with pongo2
	// after this call returns. Detect the powered marker and forward it as
	// PoweredBlock; otherwise it is plain HTML.
	if pr, ok := blocks.DecodePoweredBlock(out); ok {
		dataJSON, err := json.Marshal(pr.Data)
		if err != nil {
			return nil, internalError("marshal powered block data: " + err.Error())
		}
		return &abiv1.RenderBlockResponse{
			Powered: &abiv1.PoweredBlock{Template: pr.Template, DataJson: dataJSON},
		}, nil
	}
	return &abiv1.RenderBlockResponse{Html: out}, nil
}

// renderTag handles HOOK_RENDER_TAG: the host engine hit a plugin-declared
// tag while rendering a powered block and calls back to run it. Re-entrancy-
// free — the originating RENDER_BLOCK has already returned, so this is a fresh
// invoke. A tag fn error surfaces in RenderTagResponse.error (not an AbiError,
// which stays reserved for decode/dispatch/panic failures).
func (g *guest) renderTag(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.RenderTagRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("RenderTagRequest", err)
	}
	fn, ok := blocks.LookupTag(req.GetTagName())
	if !ok {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "unknown tag " + req.GetTagName(),
		}
	}
	ctx = contextFromRenderContext(ctx, req.GetRenderContext())
	args := unmarshalMap(req.GetArgsJson())
	if args == nil {
		args = map[string]any{}
	}
	html, err := fn(args, ctx)
	if err != nil {
		return &abiv1.RenderTagResponse{Error: err.Error()}, nil
	}
	return &abiv1.RenderTagResponse{Html: html}, nil
}

// applyFilter handles HOOK_APPLY_FILTER: run a plugin-declared filter. Same
// re-entrancy-free callback model as renderTag.
func (g *guest) applyFilter(_ context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.ApplyFilterRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("ApplyFilterRequest", err)
	}
	fn, ok := blocks.LookupFilter(req.GetFilterName())
	if !ok {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "unknown filter " + req.GetFilterName(),
		}
	}
	args := unmarshalMap(req.GetArgsJson())
	if args == nil {
		args = map[string]any{}
	}
	out, err := fn(req.GetInput(), args)
	if err != nil {
		return &abiv1.ApplyFilterResponse{Error: err.Error()}, nil
	}
	return &abiv1.ApplyFilterResponse{Output: out}, nil
}

func (g *guest) renderTemplate(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.RenderTemplateRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("RenderTemplateRequest", err)
	}
	fn, ok := g.templates.lookup(req.GetTemplateKey())
	if !ok {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "unknown template key " + req.GetTemplateKey(),
		}
	}
	ctx = contextFromRenderContext(ctx, req.GetRenderContext())
	doc := unmarshalMap(req.GetDocJson())
	if doc == nil {
		doc = map[string]any{}
	}
	component := fn(ctx, doc)
	if component == nil {
		return nil, internalError("template " + req.GetTemplateKey() + " returned nil component")
	}
	var buf bytes.Buffer
	if err := component.Render(ctx, &buf); err != nil {
		return nil, internalError("render template " + req.GetTemplateKey() + ": " + err.Error())
	}
	return &abiv1.RenderTemplateResponse{Html: buf.Bytes()}, nil
}

func (g *guest) handleHTTP(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	if g.reg.HTTPHandler == nil {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "plugin has no HTTP handler",
		}
	}
	req := &abiv1.HttpRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("HttpRequest", err)
	}
	if !g.httpHandlerInit {
		g.httpHandler = g.reg.HTTPHandler(g.services)
		g.httpHandlerInit = true
	}
	if g.httpHandler == nil {
		return nil, internalError("HTTPHandler returned nil handler")
	}

	httpReq, err := http.NewRequestWithContext(ctx, req.GetMethod(),
		(&url.URL{Path: req.GetPath(), RawQuery: req.GetRawQuery()}).String(),
		bytes.NewReader(req.GetBody()))
	if err != nil {
		return nil, internalError("build http request: " + err.Error())
	}
	for name, hv := range req.GetHeaders() {
		for _, v := range hv.GetValues() {
			httpReq.Header.Add(name, v)
		}
	}

	rec := httptest.NewRecorder()
	g.httpHandler.ServeHTTP(rec, httpReq)
	res := rec.Result()

	out := &abiv1.HttpResponse{
		Status:  int32(res.StatusCode),
		Headers: make(map[string]*abiv1.HeaderValues, len(res.Header)),
		Body:    rec.Body.Bytes(),
	}
	for name, values := range res.Header {
		out.Headers[name] = &abiv1.HeaderValues{Values: values}
	}
	return out, nil
}

func (g *guest) runJob(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.JobRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("JobRequest", err)
	}
	if !g.jobHandlersInit {
		if g.reg.JobHandlers != nil {
			g.jobHandlers = g.reg.JobHandlers(g.services)
		}
		g.jobHandlersInit = true
	}
	handler, ok := g.jobHandlers[req.GetJobType()]
	if !ok {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "unknown job type " + req.GetJobType(),
		}
	}
	// Job progress reporting needs a jobs.progress host function — an ABI
	// open item (core/docs/wasm-abi.md); v1 discards progress updates.
	progress := func(current, total int, message string) {}
	result, err := handler(ctx, req.GetConfigJson(), progress)
	if err != nil {
		return nil, internalError(err.Error())
	}
	return &abiv1.JobResponse{ResultJson: result}, nil
}

func (g *guest) load(payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.LoadRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("LoadRequest", err)
	}
	if hc := req.GetHostConfig(); hc != nil {
		g.services.AppURL = hc.GetAppUrl()
		g.services.MediaPath = hc.GetMediaPath()
	}
	if g.reg.Load != nil {
		if err := g.reg.Load(g.services); err != nil {
			return nil, internalError(err.Error())
		}
	}
	return &abiv1.LoadResponse{}, nil
}

func (g *guest) unload(ctx context.Context) (proto.Message, *abiv1.AbiError) {
	if g.reg.Unload != nil {
		if err := g.reg.Unload(ctx); err != nil {
			return nil, internalError(err.Error())
		}
	}
	return &abiv1.UnloadResponse{}, nil
}

func (g *guest) ragFetch(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	req := &abiv1.RagFetchRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("RagFetchRequest", err)
	}
	fetcher, ok := g.ragFetcher(req.GetContentType())
	if !ok {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "no content fetcher for type " + req.GetContentType(),
		}
	}
	contentID, err := uuid.Parse(req.GetContentId())
	if err != nil {
		return nil, decodeError("RagFetchRequest.content_id", err)
	}
	title, text, err := fetcher(ctx, contentID)
	if err != nil {
		return nil, internalError(err.Error())
	}
	return &abiv1.RagFetchResponse{Title: title, Text: text}, nil
}

func (g *guest) mediaHook(ctx context.Context, payload []byte) (proto.Message, *abiv1.AbiError) {
	if g.reg.MediaHooks == nil {
		return nil, &abiv1.AbiError{
			Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_UNIMPLEMENTED,
			Message: "plugin has no media hooks",
		}
	}
	req := &abiv1.MediaHookRequest{}
	if err := proto.Unmarshal(payload, req); err != nil {
		return nil, decodeError("MediaHookRequest", err)
	}
	var err error
	switch ev := req.GetEvent().(type) {
	case *abiv1.MediaHookRequest_MediaAnalyzed:
		e := ev.MediaAnalyzed
		err = g.reg.MediaHooks.OnMediaAnalyzed(ctx, plugin.MediaAnalyzedEvent{
			MediaID:      parseUUID(e.GetMediaId()),
			AnalysisID:   parseUUID(e.GetAnalysisId()),
			ContentHash:  e.GetContentHash(),
			Status:       e.GetStatus(),
			SourcePlugin: e.GetSourcePlugin(),
			SourceType:   e.GetSourceType(),
			SourceRefID:  parseUUID(e.GetSourceRefId()),
			SafeAdult:    e.GetSafeAdult(),
			SafeViolence: e.GetSafeViolence(),
			SafeRacy:     e.GetSafeRacy(),
		})
	case *abiv1.MediaHookRequest_ModerationDecision:
		e := ev.ModerationDecision
		err = g.reg.MediaHooks.OnModerationDecision(ctx, plugin.ModerationDecisionEvent{
			MediaID:        parseUUID(e.GetMediaId()),
			AnalysisID:     parseUUID(e.GetAnalysisId()),
			Status:         e.GetStatus(),
			PreviousStatus: e.GetPreviousStatus(),
			SourcePlugin:   e.GetSourcePlugin(),
			SourceType:     e.GetSourceType(),
			SourceRefID:    parseUUID(e.GetSourceRefId()),
			ModeratedBy:    parseUUID(e.GetModeratedBy()),
			Note:           e.GetNote(),
		})
	default:
		return nil, decodeError("MediaHookRequest.event", fmt.Errorf("empty oneof"))
	}
	if err != nil {
		return nil, internalError(err.Error())
	}
	return &abiv1.MediaHookResponse{}, nil
}

// --- error helpers ---

func internalError(msg string) *abiv1.AbiError {
	return &abiv1.AbiError{Code: abiv1.AbiErrorCode_ABI_ERROR_CODE_INTERNAL, Message: msg}
}

func decodeError(what string, err error) *abiv1.AbiError {
	return &abiv1.AbiError{
		Code:    abiv1.AbiErrorCode_ABI_ERROR_CODE_DECODE,
		Message: "decode " + what + ": " + err.Error(),
	}
}

// errorResponse builds a serialized error-only InvokeResponse. Marshaling a
// flat AbiError cannot fail; if the runtime is corrupted enough that it
// does, returning nil makes bn_invoke return packed 0, which the host treats
// as INTERNAL + instance discard (per wasm-abi.md).
func errorResponse(code abiv1.AbiErrorCode, msg string) []byte {
	b, err := proto.Marshal(&abiv1.InvokeResponse{Error: &abiv1.AbiError{Code: code, Message: msg}})
	if err != nil {
		return nil
	}
	return b
}

// ragFetcher returns the content fetcher the plugin registered for a content
// type (via RAGService.RegisterContentFetcher during LOAD), for
// HOOK_RAG_FETCH dispatch. Nil-safe when the guest has no RAG stub.
func (g *guest) ragFetcher(contentType string) (plugin.ContentFetcher, bool) {
	if g.rag == nil {
		return nil, false
	}
	return g.rag.Fetcher(contentType)
}

// sortedKeys returns the map's keys in sorted order (deterministic manifests).
func sortedKeys[V any](m map[string]V) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}