package adk

import (
	"context"
	"encoding/json"
	"testing"

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

func TestConstraintCheckTool(t *testing.T) {
	// Mock client that returns constraint assertions
	client := &mockClient{
		queryFunc: func(ctx context.Context, params steme.QueryParams) (*steme.QueryResult, error) {
			lifecycle := steme.LifecycleApproved
			return &steme.QueryResult{
				Assertions: []steme.AssertionResponse{
					{
						Hash:       "constraint1",
						Subject:    "auth_jwt",
						Predicate:  "must_use",
						Object:     steme.NewTextValue("Ed25519"),
						Confidence: 1.0,
						SourceHash: "source",
						Lifecycle:  lifecycle,
					},
					{
						Hash:       "constraint2",
						Subject:    "auth_jwt",
						Predicate:  "forbidden",
						Object:     steme.NewTextValue("MD5"),
						Confidence: 1.0,
						SourceHash: "source",
						Lifecycle:  lifecycle,
					},
				},
				TotalCount: 2,
			}, nil
		},
	}

	tool := NewConstraintCheckTool(client)

	input := ConstraintCheckInput{
		Context: "auth_jwt",
	}

	inputBytes, _ := json.Marshal(input)
	outputBytes, err := tool.Execute(context.Background(), inputBytes)
	if err != nil {
		t.Fatalf("constraint check failed: %v", err)
	}

	var output ConstraintCheckOutput
	if err := json.Unmarshal(outputBytes, &output); err != nil {
		t.Fatalf("failed to unmarshal output: %v", err)
	}

	// Verify constraints were returned
	if len(output.Constraints) != 2 {
		t.Errorf("expected 2 constraints, got %d", len(output.Constraints))
	}

	// Check for must-use constraint
	foundMustUse := false
	for _, c := range output.Constraints {
		if c.MustUse == "Ed25519" {
			foundMustUse = true
		}
	}
	if !foundMustUse {
		t.Error("expected to find Ed25519 must-use constraint")
	}

	// Check for forbidden constraint
	foundForbidden := false
	for _, c := range output.Constraints {
		if c.Forbidden == "MD5" {
			foundForbidden = true
		}
	}
	if !foundForbidden {
		t.Error("expected to find MD5 forbidden constraint")
	}
}

func TestTraceTool(t *testing.T) {
	// Mock client that returns audit records
	agentID := "deadbeef00000000000000000000000000000000000000000000000000000000"
	subject := "Tesla_Inc"
	predicate := "has_revenue"
	resultHash := "result_hash_123"
	lens := steme.LensConsensus

	client := &mockClient{
		traceFunc: func(ctx context.Context, params steme.TraceParams) (*steme.TraceResult, error) {
			return &steme.TraceResult{
				Audits: []steme.QueryAuditRecord{
					{
						QueryID:   "query_123",
						AgentID:   &agentID,
						Timestamp: 1704067200,
						Params: steme.QueryParamsAudit{
							Subject:   &subject,
							Predicate: &predicate,
							Lens:      &lens,
						},
						ResultHash:       &resultHash,
						ResultConfidence: 0.95,
						ContributingAssertions: []steme.ContributingAssertion{
							{
								AssertionHash: "assertion_1",
								Weight:        1.0,
								SourceHash:    "source_1",
								Lifecycle:     steme.LifecycleApproved,
							},
							{
								AssertionHash: "assertion_2",
								Weight:        0.5,
								SourceHash:    "source_2",
								Lifecycle:     steme.LifecycleProposed,
							},
						},
					},
				},
				TotalCount:    1,
				AgentID:       agentID,
				FromTimestamp: 1704067200,
				ToTimestamp:   1704153600,
			}, nil
		},
	}

	tool := NewTraceTool(client)

	input := TraceInput{
		AgentID: agentID,
		From:    "1704067200",
		To:      "1704153600",
		Subject: "Tesla*",
	}

	inputBytes, _ := json.Marshal(input)
	outputBytes, err := tool.Execute(context.Background(), inputBytes)
	if err != nil {
		t.Fatalf("trace failed: %v", err)
	}

	var output TraceOutput
	if err := json.Unmarshal(outputBytes, &output); err != nil {
		t.Fatalf("failed to unmarshal output: %v", err)
	}

	// Verify output
	if len(output.Queries) != 1 {
		t.Errorf("expected 1 query, got %d", len(output.Queries))
	}

	if len(output.Queries) > 0 {
		query := output.Queries[0]

		if query.QueryID != "query_123" {
			t.Errorf("expected query_id query_123, got %s", query.QueryID)
		}

		if query.Subject != "Tesla_Inc" {
			t.Errorf("expected subject Tesla_Inc, got %s", query.Subject)
		}

		if query.Predicate != "has_revenue" {
			t.Errorf("expected predicate has_revenue, got %s", query.Predicate)
		}

		if query.Lens != "consensus" {
			t.Errorf("expected lens consensus, got %s", query.Lens)
		}

		if query.Confidence != 0.95 {
			t.Errorf("expected confidence 0.95, got %f", query.Confidence)
		}

		if len(query.Contributing) != 2 {
			t.Errorf("expected 2 contributing assertions, got %d", len(query.Contributing))
		}
	}
}