package mcp

import (
	"context"

	"github.com/yan1h/agent-coding-workflow/internal/audit"
)

type mcpTokenIDCtxKey struct{}

func withMcpTokenID(ctx context.Context, tokenID string) context.Context {
	return context.WithValue(ctx, mcpTokenIDCtxKey{}, tokenID)
}

func mcpTokenIDFromCtx(ctx context.Context) string {
	if v, ok := ctx.Value(mcpTokenIDCtxKey{}).(string); ok {
		return v
	}
	return ""
}

type wrappedRecorder struct {
	inner audit.Recorder
}

// AuditWrap wraps an audit.Recorder so that every Record call automatically
// injects "via_mcp" and "mcp_token_id" into the entry's Metadata when the
// context carries an MCP token ID. Returns nil when inner is nil.
func AuditWrap(inner audit.Recorder) audit.Recorder {
	if inner == nil {
		return nil
	}
	return &wrappedRecorder{inner: inner}
}

func (r *wrappedRecorder) Record(ctx context.Context, e audit.Entry) error {
	if tid := mcpTokenIDFromCtx(ctx); tid != "" {
		if e.Metadata == nil {
			e.Metadata = map[string]any{}
		}
		e.Metadata["via_mcp"] = true
		e.Metadata["mcp_token_id"] = tid
	}
	return r.inner.Record(ctx, e)
}