stemedb/sdk/go/adk/tools.go

// Package adk provides ADK-Go tool wrappers for StemeDB.
//
// This package wraps the base StemeDB SDK with tool definitions suitable
// for Google's Agent Development Kit (ADK-Go). It provides:
//
//   - QueryTool - Query with lens-based resolution
//   - AssertTool - Create assertions with confidence
//   - ConstraintCheckTool - Pre-flight validation
//   - TraceTool - Audit trail queries
//   - SupersedeTool - Epoch/correction management
//
// Since ADK-Go is not yet publicly released, these tools use interface-based
// design and can be adapted to any ADK implementation.
//
// Example usage:
//
//	client := steme.NewClient("http://localhost:3000", signer)
//	tools := adk.AllTools(client)
//
//	// Use tools with your ADK agent configuration
//	agent := configureAgent(tools)
package adk

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/orchard9/stemedb-go/steme"
)

// Tool represents a generic ADK tool interface.
//
// This interface is designed to be compatible with any ADK implementation.
type Tool interface {
	Name() string
	Description() string
	Execute(ctx context.Context, input []byte) ([]byte, error)
}

// EpistemeClient defines the interface for StemeDB operations.
//
// This interface allows for easy mocking and testing.
type EpistemeClient interface {
	Query(ctx context.Context, params steme.QueryParams) (*steme.QueryResult, error)
	Assert(ctx context.Context, assertion steme.Assertion) (string, error)
	Trace(ctx context.Context, params steme.TraceParams) (*steme.TraceResult, error)
	Supersede(ctx context.Context, params steme.SupersedeParams) (*steme.SupersedeResult, error)
}

// QueryTool provides lens-based knowledge retrieval.
//
// Used by all agents to retrieve current knowledge with conflict resolution.
type QueryTool struct {
	client EpistemeClient
}

// NewQueryTool creates a new Query tool.
func NewQueryTool(client EpistemeClient) *QueryTool {
	return &QueryTool{client: client}
}

// Name returns the tool name.
func (t *QueryTool) Name() string {
	return "episteme_query"
}

// Description returns the tool description.
func (t *QueryTool) Description() string {
	return "Query Episteme knowledge graph with lens-based conflict resolution. " +
		"Returns resolved value with confidence score and source provenance. " +
		"CRITICAL: Always filter by lifecycle=approved for production decisions."
}

// Execute performs the query operation.
func (t *QueryTool) Execute(ctx context.Context, input []byte) ([]byte, error) {
	var params QueryInput
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("invalid query input: %w", err)
	}

	// Build StemeDB query params
	builder := steme.NewQuery()

	if params.Subject != "" {
		builder.WithSubject(params.Subject)
	}

	if params.Predicate != "" {
		builder.WithPredicate(params.Predicate)
	}

	if params.Lens != "" {
		lens := parseLen(params.Lens)
		builder.WithLens(lens)
	}

	if params.Lifecycle != "" {
		lifecycle := parseLifecycle(params.Lifecycle)
		builder.WithLifecycle(lifecycle)
	}

	// Execute query
	result, err := t.client.Query(ctx, builder.Build())
	if err != nil {
		errorOutput, marshalErr := json.Marshal(QueryOutput{
			Error: err.Error(),
		})
		if marshalErr != nil {
			return nil, fmt.Errorf("query failed: %w", err)
		}
		return errorOutput, nil
	}

	// Convert to QueryOutput format
	output := convertQueryResult(result, params.MinConfidence)

	outputBytes, err := json.Marshal(output)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal query output: %w", err)
	}
	return outputBytes, nil
}

// AssertTool stores knowledge with lifecycle and confidence.
//
// Used by Research Agent and Human Supervisor to store assertions.
type AssertTool struct {
	client EpistemeClient
}

// NewAssertTool creates a new Assert tool.
func NewAssertTool(client EpistemeClient) *AssertTool {
	return &AssertTool{client: client}
}

// Name returns the tool name.
func (t *AssertTool) Name() string {
	return "episteme_assert"
}

// Description returns the tool description.
func (t *AssertTool) Description() string {
	return "Store knowledge assertion in Episteme. " +
		"Always include source_hash for provenance. " +
		"Use confidence to express uncertainty. " +
		"Mark as lifecycle=proposed unless explicitly approved."
}

// Execute performs the assert operation.
func (t *AssertTool) Execute(ctx context.Context, input []byte) ([]byte, error) {
	var params AssertInput
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("invalid assert input: %w", err)
	}

	// Build assertion
	builder := steme.NewAssertion(params.Subject, params.Predicate)

	// Set object value based on type
	if err := setObjectValue(builder, params.Object); err != nil {
		errorOutput, marshalErr := json.Marshal(AssertOutput{
			Success: false,
			Error:   err.Error(),
		})
		if marshalErr != nil {
			return nil, fmt.Errorf("invalid object value: %w", err)
		}
		return errorOutput, nil
	}

	builder.WithConfidence(float64(params.Confidence)).
		WithSourceHash(params.SourceHash)

	if params.Lifecycle != "" {
		lifecycle := parseLifecycle(params.Lifecycle)
		builder.WithLifecycle(lifecycle)
	}

	if params.ParentHash != "" {
		builder.WithParentHash(params.ParentHash)
	}

	// Execute assertion
	hash, err := t.client.Assert(ctx, builder.Build())
	if err != nil {
		errorOutput, marshalErr := json.Marshal(AssertOutput{
			Success: false,
			Error:   err.Error(),
		})
		if marshalErr != nil {
			return nil, fmt.Errorf("assertion failed: %w", err)
		}
		return errorOutput, nil
	}

	outputBytes, err := json.Marshal(AssertOutput{
		Hash:    hash,
		Success: true,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to marshal assert output: %w", err)
	}
	return outputBytes, nil
}

// ConstraintCheckTool provides pre-flight validation for Implementation Agent.
//
// Returns must-use and forbidden patterns for a given context.
type ConstraintCheckTool struct {
	client EpistemeClient
}

// NewConstraintCheckTool creates a new ConstraintCheck tool.
func NewConstraintCheckTool(client EpistemeClient) *ConstraintCheckTool {
	return &ConstraintCheckTool{client: client}
}

// Name returns the tool name.
func (t *ConstraintCheckTool) Name() string {
	return "episteme_constraint_check"
}

// Description returns the tool description.
func (t *ConstraintCheckTool) Description() string {
	return "Check for must-use and forbidden patterns before code generation. " +
		"Returns constraints with explanations for contrastive learning. " +
		"Implementation Agent MUST call this before writing code."
}

// Execute performs the constraint check.
func (t *ConstraintCheckTool) Execute(ctx context.Context, input []byte) ([]byte, error) {
	var params ConstraintCheckInput
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("invalid constraint check input: %w", err)
	}

	// Query for constraints in the given context
	// This queries for assertions with predicate "must_use" or "forbidden"
	builder := steme.NewQuery().
		WithSubject(params.Context).
		WithLifecycle(steme.LifecycleApproved). // Only approved constraints
		WithLens(steme.LensAuthority)

	result, err := t.client.Query(ctx, builder.Build())
	if err != nil {
		return nil, fmt.Errorf("constraint query failed: %w", err)
	}

	// Convert assertions to constraints
	constraints := convertToConstraints(result)

	outputBytes, err := json.Marshal(ConstraintCheckOutput{
		Constraints: constraints,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to marshal constraint output: %w", err)
	}
	return outputBytes, nil
}

// TraceTool provides audit trail queries for incident investigation.
//
// Used by On-Call SRE and Human Supervisor to trace agent decisions.
type TraceTool struct {
	client EpistemeClient
}

// NewTraceTool creates a new Trace tool.
func NewTraceTool(client EpistemeClient) *TraceTool {
	return &TraceTool{client: client}
}

// Name returns the tool name.
func (t *TraceTool) Name() string {
	return "episteme_trace"
}

// Description returns the tool description.
func (t *TraceTool) Description() string {
	return "Trace agent queries for incident investigation. " +
		"Shows what an agent queried, when, and what values were returned. " +
		"CRITICAL for debugging production incidents."
}

// Execute performs the trace operation.
func (t *TraceTool) Execute(ctx context.Context, input []byte) ([]byte, error) {
	var params TraceInput
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("invalid trace input: %w", err)
	}

	// Build StemeDB trace params
	traceParams := steme.TraceParams{
		AgentID: params.AgentID,
		From:    params.From,
		To:      params.To,
		Subject: params.Subject,
		Limit:   100, // Default limit
	}

	// Execute trace
	result, err := t.client.Trace(ctx, traceParams)
	if err != nil {
		return nil, fmt.Errorf("trace query failed: %w", err)
	}

	// Convert to TraceOutput format
	output := convertTraceResult(result)

	outputBytes, err := json.Marshal(output)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal trace output: %w", err)
	}
	return outputBytes, nil
}

// SupersedeTool provides error correction with cascade tracking.
//
// Used by Human Supervisor to correct bad assertions.
type SupersedeTool struct {
	client EpistemeClient
}

// NewSupersedeTool creates a new Supersede tool.
func NewSupersedeTool(client EpistemeClient) *SupersedeTool {
	return &SupersedeTool{client: client}
}

// Name returns the tool name.
func (t *SupersedeTool) Name() string {
	return "episteme_supersede"
}

// Description returns the tool description.
func (t *SupersedeTool) Description() string {
	return "Supersede a bad assertion with correction tracking. " +
		"Shows cascade impact on downstream assertions. " +
		"Used for incident remediation and knowledge correction."
}

// Execute performs the supersede operation.
func (t *SupersedeTool) Execute(ctx context.Context, input []byte) ([]byte, error) {
	var params SupersedeInput
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("invalid supersede input: %w", err)
	}

	// Convert SupersessionType
	superType := parseSupersessionType(params.Type)

	// Build SupersedeParams
	supersedeParams := steme.SupersedeParams{
		TargetHash:       params.Hash,
		SupersessionType: superType,
		Reason:           params.Reason,
		NewHash:          params.NewHash,
		AgentID:          params.AgentID,
		Signature:        params.Signature,
	}

	// Execute supersede
	result, err := t.client.Supersede(ctx, supersedeParams)
	if err != nil {
		return nil, fmt.Errorf("supersede failed: %w", err)
	}

	// Convert to SupersedeOutput format
	output := SupersedeOutput{
		Success:          true,
		TargetHash:       result.TargetHash,
		SupersessionType: string(result.SupersessionType),
		Timestamp:        result.Timestamp,
	}

	outputBytes, err := json.Marshal(output)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal supersede output: %w", err)
	}
	return outputBytes, nil
}

// AllTools returns all available Episteme tools.
//
// Use this for agent configuration:
//
//	tools := adk.AllTools(client)
func AllTools(client EpistemeClient) []Tool {
	return []Tool{
		NewQueryTool(client),
		NewAssertTool(client),
		NewConstraintCheckTool(client),
		NewTraceTool(client),
		NewSupersedeTool(client),
	}
}

// Helper functions

func parseLen(lens string) steme.Lens {
	switch lens {
	case "recency":
		return steme.LensRecency
	case "consensus":
		return steme.LensConsensus
	case "authority":
		return steme.LensAuthority
	case "vote_aware_consensus":
		return steme.LensVoteAwareConsensus
	case "trust_aware_authority":
		return steme.LensTrustAwareAuthority
	default:
		return steme.LensConsensus // Default to consensus
	}
}

func parseLifecycle(lifecycle string) steme.LifecycleStage {
	switch lifecycle {
	case "proposed":
		return steme.LifecycleProposed
	case "under_review":
		return steme.LifecycleUnderReview
	case "approved":
		return steme.LifecycleApproved
	case "deprecated":
		return steme.LifecycleDeprecated
	case "rejected":
		return steme.LifecycleRejected
	default:
		return steme.LifecycleProposed // Default to proposed
	}
}

func parseSupersessionType(t string) steme.SupersessionType {
	switch t {
	case "Invalidate", "invalidate":
		return steme.SupersessionInvalidate
	case "Temporal", "temporal":
		return steme.SupersessionTemporal
	case "Refinement", "refinement":
		return steme.SupersessionRefinement
	case "RequiresReview", "requires_review":
		return steme.SupersessionRequiresReview
	case "Additive", "additive":
		return steme.SupersessionAdditive
	default:
		return steme.SupersessionInvalidate // Default to invalidate
	}
}

func setObjectValue(builder *steme.AssertionBuilder, obj any) error {
	switch v := obj.(type) {
	case string:
		builder.WithText(v)
	case float64:
		builder.WithNumber(v)
	case bool:
		builder.WithBoolean(v)
	case map[string]any:
		// Handle ObjectValue from JSON
		if typeStr, ok := v["type"].(string); ok {
			if value, ok := v["value"]; ok {
				switch typeStr {
				case "Text":
					if str, ok := value.(string); ok {
						builder.WithText(str)
					}
				case "Number":
					if num, ok := value.(float64); ok {
						builder.WithNumber(num)
					}
				case "Boolean":
					if b, ok := value.(bool); ok {
						builder.WithBoolean(b)
					}
				case "Reference":
					if ref, ok := value.(string); ok {
						builder.WithReference(ref)
					}
				}
			}
		}
	default:
		return fmt.Errorf("unsupported object type: %T", obj)
	}
	return nil
}

func convertQueryResult(result *steme.QueryResult, minConfidence float32) QueryOutput {
	if len(result.Assertions) == 0 {
		return QueryOutput{
			Error: "no results found",
		}
	}

	// Take the first result (lens already applied resolution)
	assertion := result.Assertions[0]

	// Check confidence threshold
	confidence := float32(assertion.Confidence)
	if minConfidence > 0 && confidence < minConfidence {
		return QueryOutput{
			Value:      assertion.Object.Value,
			Confidence: confidence,
			Lifecycle:  string(assertion.Lifecycle),
			QueryID:    generateQueryID(),
			Error:      fmt.Sprintf("confidence %.2f below threshold %.2f", confidence, minConfidence),
		}
	}

	// Convert sources
	sources := make([]Source, 0, len(result.Assertions))
	for _, a := range result.Assertions {
		sources = append(sources, Source{
			Hash:       a.Hash,
			SourceHash: a.SourceHash,
			Weight:     float32(a.Confidence),
		})
	}

	return QueryOutput{
		Value:      assertion.Object.Value,
		Confidence: confidence,
		Lifecycle:  string(assertion.Lifecycle),
		Sources:    sources,
		QueryID:    generateQueryID(),
	}
}

func convertToConstraints(result *steme.QueryResult) []Constraint {
	constraints := make([]Constraint, 0)

	for _, assertion := range result.Assertions {
		constraint := Constraint{
			Subject: assertion.Subject,
		}

		// Determine if this is a must-use or forbidden constraint
		// based on predicate
		switch assertion.Predicate {
		case "must_use":
			if text, ok := assertion.Object.Value.(string); ok {
				constraint.MustUse = text
			}
		case "forbidden":
			if text, ok := assertion.Object.Value.(string); ok {
				constraint.Forbidden = text
			}
		}

		// Extract reason from parent assertion if available
		// (This is a simplified implementation)
		constraint.Reason = "Approved organizational pattern"

		constraints = append(constraints, constraint)
	}

	return constraints
}

func generateQueryID() string {
	// Generate a unique query ID for audit trail
	// In production, this would be generated server-side
	return fmt.Sprintf("query_%d", time.Now().UnixNano())
}

func convertTraceResult(result *steme.TraceResult) TraceOutput {
	queries := make([]QueryTrace, 0, len(result.Audits))

	for _, audit := range result.Audits {
		// Extract subject and predicate from params
		subject := ""
		if audit.Params.Subject != nil {
			subject = *audit.Params.Subject
		}

		predicate := ""
		if audit.Params.Predicate != nil {
			predicate = *audit.Params.Predicate
		}

		// Extract lens (convert to lowercase for consistency)
		lens := "consensus" // Default
		if audit.Params.Lens != nil {
			lensVal := string(*audit.Params.Lens)
			// Convert "Consensus" -> "consensus", etc.
			switch lensVal {
			case "Recency":
				lens = "recency"
			case "Consensus":
				lens = "consensus"
			case "Authority":
				lens = "authority"
			case "VoteAwareConsensus":
				lens = "vote_aware_consensus"
			case "TrustAwareAuthority":
				lens = "trust_aware_authority"
			default:
				lens = "consensus"
			}
		}

		// Extract result value (simplified - just use hash if available)
		resultValue := ""
		if audit.ResultHash != nil {
			resultValue = *audit.ResultHash
		}

		// Extract contributing assertion hashes
		contributing := make([]string, 0, len(audit.ContributingAssertions))
		for _, ca := range audit.ContributingAssertions {
			contributing = append(contributing, ca.AssertionHash)
		}

		queries = append(queries, QueryTrace{
			QueryID:                audit.QueryID,
			Timestamp:              fmt.Sprintf("%d", audit.Timestamp),
			Subject:                subject,
			Predicate:              predicate,
			Lens:                   lens,
			Result:                 resultValue,
			Confidence:             audit.ResultConfidence,
			Contributing:           contributing,
		})
	}

	return TraceOutput{
		Queries: queries,
	}
}