core/plugin/wasmguest/dispatch.go

package wasmguest

import (
	"bytes"
	"context"
	"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/plugin"
	"github.com/google/uuid"
	"google.golang.org/protobuf/proto"
)

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

// 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). The capability-stub
	// WO populates its interface fields with host_call-backed
	// implementations; LOAD fills AppURL/MediaPath from HostConfig.
	services plugin.CoreServices

	// ragFetchers records RAGService.RegisterContentFetcher calls made
	// during LOAD, dispatched via HOOK_RAG_FETCH.
	ragFetchers map[string]plugin.ContentFetcher

	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(),
		ragFetchers: make(map[string]plugin.ContentFetcher),
	}
	g.services.RAGService = &guestRAGService{fetchers: g.ragFetchers}
	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)
		}
	}()
	// 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_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{}
	}
	return &abiv1.RenderBlockResponse{Html: fn(ctx, content)}, 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.ragFetchers[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
}

// --- guest-side RAG service (fetcher capture) ---

// guestRAGService records content-fetcher registrations so HOOK_RAG_FETCH
// can dispatch to them. Query/OnContentChanged become host capability calls
// in the capability-stub WO; until then they report unimplemented.
type guestRAGService struct {
	fetchers map[string]plugin.ContentFetcher
}

func (s *guestRAGService) RegisterContentFetcher(contentType string, fetcher plugin.ContentFetcher) {
	s.fetchers[contentType] = fetcher
}

func (s *guestRAGService) Query(context.Context, string, int) ([]plugin.RAGResult, error) {
	return nil, fmt.Errorf("wasmguest: rag.query host capability not wired yet")
}

func (s *guestRAGService) OnContentChanged(context.Context, string, uuid.UUID) {}

// 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
}