package integration

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/anthonyrawlins/bzzz/pkg/dht"
	"github.com/anthonyrawlins/bzzz/pkg/ucxl"
)

// Phase 1 Integration Tests for BZZZ-RUSTLE Mock Implementation
// These tests validate that mock components work together as designed

func TestPhase1MockIntegration(t *testing.T) {
	ctx := context.Background()
	
	t.Run("MockDHT_Basic_Operations", func(t *testing.T) {
		testMockDHTBasicOperations(t, ctx)
	})
	
	t.Run("UCXL_Address_Consistency", func(t *testing.T) {
		testUCXLAddressConsistency(t)
	})
	
	t.Run("MockDHT_UCXL_Integration", func(t *testing.T) {
		testMockDHTUCXLIntegration(t, ctx)
	})
	
	t.Run("Cross_Language_Compatibility", func(t *testing.T) {
		testCrossLanguageCompatibility(t)
	})
}

func testMockDHTBasicOperations(t *testing.T, ctx context.Context) {
	// Test that Mock DHT provides same interface as real DHT
	mockDHT := dht.NewMockDHT()
	
	// Test storage operations
	testKey := "ucxl://coordinator.local:config@bzzz:cluster/bootstrap"
	testValue := []byte(`{
		"cluster_id": "bzzz-test-cluster",
		"bootstrap_nodes": ["192.168.1.100:8080"],
		"admin_key_threshold": 3,
		"total_admin_keys": 5
	}`)
	
	// Store configuration
	err := mockDHT.PutValue(ctx, testKey, testValue)
	if err != nil {
		t.Fatalf("Failed to store value in mock DHT: %v", err)
	}
	
	// Retrieve configuration
	retrieved, err := mockDHT.GetValue(ctx, testKey)
	if err != nil {
		t.Fatalf("Failed to retrieve value from mock DHT: %v", err)
	}
	
	if string(retrieved) != string(testValue) {
		t.Fatalf("Retrieved value doesn't match stored value")
	}
	
	// Test provider announcement (for service discovery)
	providerId := "rustle-browser-001"
	err = mockDHT.Provide(ctx, testKey, providerId)
	if err != nil {
		t.Fatalf("Failed to announce provider: %v", err)
	}
	
	// Find providers
	providers, err := mockDHT.FindProviders(ctx, testKey)
	if err != nil {
		t.Fatalf("Failed to find providers: %v", err)
	}
	
	found := false
	for _, provider := range providers {
		if provider == providerId {
			found = true
			break
		}
	}
	
	if !found {
		t.Fatalf("Provider %s not found in provider list", providerId)
	}
	
	t.Logf("✓ Mock DHT: Basic operations working correctly")
}

func testUCXLAddressConsistency(t *testing.T) {
	// Test that UCXL addresses work consistently across different use cases
	testCases := []struct {
		name        string
		agent       string
		role        string
		project     string
		task        string
		path        string
		temporal    string
		shouldMatch bool
	}{
		{
			name:        "Coordinator Config",
			agent:       "coordinator-001",
			role:        "leader",
			project:     "bzzz", 
			task:        "config",
			path:        "/cluster/bootstrap",
			temporal:    "",
			shouldMatch: true,
		},
		{
			name:        "RUSTLE Browser Request",
			agent:       "rustle-browser",
			role:        "client",
			project:     "bzzz",
			task:        "query",
			path:        "/models/available",
			temporal:    "^/",
			shouldMatch: true,
		},
		{
			name:        "Wildcard Search",
			agent:       "*",
			role:        "*", 
			project:     "bzzz",
			task:        "*",
			path:        "/models/*",
			temporal:    "",
			shouldMatch: true,
		},
	}
	
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Generate address
			address, err := ucxl.GenerateUCXLAddress(
				tc.agent, tc.role, tc.project, tc.task, tc.path, tc.temporal,
			)
			if err != nil {
				t.Fatalf("Failed to generate UCXL address: %v", err)
			}
			
			// Parse address
			parsed, err := ucxl.ParseUCXLAddress(address)
			if err != nil {
				t.Fatalf("Failed to parse UCXL address: %v", err)
			}
			
			// Verify round-trip consistency
			regenerated, err := ucxl.GenerateUCXLAddress(
				parsed.Agent, parsed.Role, parsed.Project, parsed.Task, parsed.Path, parsed.Temporal,
			)
			if err != nil {
				t.Fatalf("Failed to regenerate UCXL address: %v", err)
			}
			
			// Parse regenerated address to ensure consistency
			_, err = ucxl.ParseUCXLAddress(regenerated)
			if err != nil {
				t.Fatalf("Failed to parse regenerated UCXL address: %v", err)
			}
			
			t.Logf("✓ UCXL Address: %s → %s", tc.name, address)
		})
	}
}

func testMockDHTUCXLIntegration(t *testing.T, ctx context.Context) {
	// Test that Mock DHT and UCXL work together for realistic scenarios
	mockDHT := dht.NewMockDHT()
	
	// Scenario 1: Store configuration using UCXL addressing
	configAddr, err := ucxl.GenerateUCXLAddress(
		"coordinator-001", "leader", "bzzz", "config", "/cluster/nodes", "",
	)
	if err != nil {
		t.Fatalf("Failed to generate config address: %v", err)
	}
	
	nodeConfig := []byte(`{
		"node_id": "bzzz-node-001",
		"address": "192.168.1.101:8080",
		"capabilities": ["storage", "processing", "coordination"],
		"vram_gb": 24,
		"model_slots": ["llama-3.1-70b", "qwen-2.5-32b"]
	}`)
	
	err = mockDHT.PutValue(ctx, configAddr, nodeConfig)
	if err != nil {
		t.Fatalf("Failed to store node config: %v", err)
	}
	
	// Scenario 2: RUSTLE browser queries for available models
	modelQueryAddr, err := ucxl.GenerateUCXLAddress(
		"rustle-browser", "client", "bzzz", "query", "/models/available", "^/",
	)
	if err != nil {
		t.Fatalf("Failed to generate model query address: %v", err)
	}
	
	// Store available models information
	modelsInfo := []byte(`{
		"available_models": [
			{
				"name": "llama-3.1-70b",
				"node": "bzzz-node-001",
				"status": "available",
				"vram_required": 20
			},
			{
				"name": "qwen-2.5-32b", 
				"node": "bzzz-node-002",
				"status": "available",
				"vram_required": 16
			}
		],
		"timestamp": "2025-01-10T15:30:00Z"
	}`)
	
	err = mockDHT.PutValue(ctx, modelQueryAddr, modelsInfo)
	if err != nil {
		t.Fatalf("Failed to store models info: %v", err)
	}
	
	// Scenario 3: Verify both configurations can be retrieved
	retrievedConfig, err := mockDHT.GetValue(ctx, configAddr)
	if err != nil {
		t.Fatalf("Failed to retrieve node config: %v", err)
	}
	
	retrievedModels, err := mockDHT.GetValue(ctx, modelQueryAddr)
	if err != nil {
		t.Fatalf("Failed to retrieve models info: %v", err)
	}
	
	if len(retrievedConfig) == 0 || len(retrievedModels) == 0 {
		t.Fatalf("Retrieved data is empty")
	}
	
	// Test stats
	stats := mockDHT.GetStats()
	if stats.TotalKeys < 2 {
		t.Fatalf("Expected at least 2 keys in storage, got %d", stats.TotalKeys)
	}
	
	t.Logf("✓ Mock DHT + UCXL Integration: Successfully stored and retrieved configurations")
	t.Logf("  - Node config address: %s", configAddr)
	t.Logf("  - Models query address: %s", modelQueryAddr) 
	t.Logf("  - Total keys in storage: %d", stats.TotalKeys)
}

func testCrossLanguageCompatibility(t *testing.T) {
	// Test compatibility patterns between Go (BZZZ) and Rust (RUSTLE)
	// This validates that both implementations follow the same addressing schemes
	
	// Test cases that should work identically in both languages
	compatibilityTests := []struct {
		name     string
		address  string
		expected map[string]string
	}{
		{
			name:    "Basic BZZZ Config",
			address: "ucxl://coordinator:leader@bzzz:config/cluster/bootstrap",
			expected: map[string]string{
				"agent":   "coordinator",
				"role":    "leader",
				"project": "bzzz",
				"task":    "config",
				"path":    "/cluster/bootstrap",
			},
		},
		{
			name:    "RUSTLE Query with Temporal",
			address: "ucxl://browser:client@bzzz:query/models/available*^/",
			expected: map[string]string{
				"agent":    "browser", 
				"role":     "client",
				"project":  "bzzz",
				"task":     "query",
				"path":     "/models/available",
				"temporal": "^/",
			},
		},
		{
			name:    "Wildcard Pattern",
			address: "ucxl://*:*@bzzz:*/*",
			expected: map[string]string{
				"agent":   "*",
				"role":    "*",
				"project": "bzzz",
				"task":    "*",
				"path":    "/",
			},
		},
	}
	
	for _, tc := range compatibilityTests {
		t.Run(tc.name, func(t *testing.T) {
			// Parse address using Go implementation
			parsed, err := ucxl.ParseUCXLAddress(tc.address)
			if err != nil {
				t.Fatalf("Failed to parse address with Go implementation: %v", err)
			}
			
			// Verify expected fields match
			if tc.expected["agent"] != "" && parsed.Agent != tc.expected["agent"] {
				t.Errorf("Agent mismatch: got %s, want %s", parsed.Agent, tc.expected["agent"])
			}
			if tc.expected["role"] != "" && parsed.Role != tc.expected["role"] {
				t.Errorf("Role mismatch: got %s, want %s", parsed.Role, tc.expected["role"])
			}
			if tc.expected["project"] != "" && parsed.Project != tc.expected["project"] {
				t.Errorf("Project mismatch: got %s, want %s", parsed.Project, tc.expected["project"])
			}
			if tc.expected["task"] != "" && parsed.Task != tc.expected["task"] {
				t.Errorf("Task mismatch: got %s, want %s", parsed.Task, tc.expected["task"])
			}
			if tc.expected["path"] != "" && parsed.Path != tc.expected["path"] {
				t.Errorf("Path mismatch: got %s, want %s", parsed.Path, tc.expected["path"])
			}
			
			t.Logf("✓ Cross-Language Compatibility: %s", tc.name)
		})
	}
}

// TestPhase1Scenarios tests realistic integration scenarios
func TestPhase1Scenarios(t *testing.T) {
	ctx := context.Background()
	
	t.Run("Scenario_Bootstrap_Cluster", func(t *testing.T) {
		testBootstrapClusterScenario(t, ctx)
	})
	
	t.Run("Scenario_RUSTLE_Model_Discovery", func(t *testing.T) {
		testRUSTLEModelDiscoveryScenario(t, ctx)
	})
}

func testBootstrapClusterScenario(t *testing.T, ctx context.Context) {
	// Simulate cluster bootstrap process using mock components
	mockDHT := dht.NewMockDHT()
	
	// Step 1: Store initial cluster configuration
	clusterConfigAddr, _ := ucxl.GenerateUCXLAddress(
		"admin", "bootstrap", "bzzz", "config", "/cluster/initial", "",
	)
	
	initialConfig := []byte(`{
		"cluster_name": "bzzz-production",
		"bootstrap_complete": false,
		"admin_nodes": ["192.168.1.100", "192.168.1.101", "192.168.1.102"],
		"required_admin_shares": 3,
		"total_admin_shares": 5
	}`)
	
	err := mockDHT.PutValue(ctx, clusterConfigAddr, initialConfig)
	if err != nil {
		t.Fatalf("Failed to store initial cluster config: %v", err)
	}
	
	// Step 2: Each admin node announces itself
	adminNodes := []string{"admin-001", "admin-002", "admin-003"}
	for i, nodeId := range adminNodes {
		nodeAddr, _ := ucxl.GenerateUCXLAddress(
			nodeId, "admin", "bzzz", "announce", "/node/ready", "",
		)
		
		nodeInfo := []byte(fmt.Sprintf(`{
			"node_id": "%s",
			"address": "192.168.1.%d:8080", 
			"public_key": "mock-key-%s",
			"ready": true,
			"timestamp": "%s"
		}`, nodeId, 100+i, nodeId, time.Now().Format(time.RFC3339)))
		
		err := mockDHT.PutValue(ctx, nodeAddr, nodeInfo)
		if err != nil {
			t.Fatalf("Failed to announce admin node %s: %v", nodeId, err)
		}
		
		// Announce as provider for admin services
		err = mockDHT.Provide(ctx, "bzzz:admin:services", nodeId)
		if err != nil {
			t.Fatalf("Failed to announce admin provider %s: %v", nodeId, err)
		}
	}
	
	// Step 3: Verify all admin nodes are discoverable
	adminProviders, err := mockDHT.FindProviders(ctx, "bzzz:admin:services")
	if err != nil {
		t.Fatalf("Failed to find admin providers: %v", err)
	}
	
	if len(adminProviders) != len(adminNodes) {
		t.Fatalf("Expected %d admin providers, got %d", len(adminNodes), len(adminProviders))
	}
	
	// Step 4: Update cluster config to indicate bootstrap completion
	updatedConfig := []byte(`{
		"cluster_name": "bzzz-production",
		"bootstrap_complete": true,
		"active_admin_nodes": 3,
		"cluster_ready": true,
		"bootstrap_timestamp": "` + time.Now().Format(time.RFC3339) + `"
	}`)
	
	err = mockDHT.PutValue(ctx, clusterConfigAddr, updatedConfig)
	if err != nil {
		t.Fatalf("Failed to update cluster config: %v", err)
	}
	
	t.Logf("✓ Bootstrap Cluster Scenario: Successfully simulated cluster bootstrap")
	t.Logf("  - Admin nodes announced: %d", len(adminNodes))
	t.Logf("  - Admin providers discoverable: %d", len(adminProviders))
}

func testRUSTLEModelDiscoveryScenario(t *testing.T, ctx context.Context) {
	// Simulate RUSTLE browser discovering available models through mock BZZZ
	mockDHT := dht.NewMockDHT()
	
	// Step 1: Processing nodes announce their capabilities
	processingNodes := []struct {
		nodeId string
		models []string
		vram   int
	}{
		{"worker-001", []string{"llama-3.1-8b", "qwen-2.5-7b"}, 12},
		{"worker-002", []string{"llama-3.1-70b"}, 80}, 
		{"worker-003", []string{"mixtral-8x7b", "qwen-2.5-32b"}, 48},
	}
	
	for i, node := range processingNodes {
		nodeAddr, _ := ucxl.GenerateUCXLAddress(
			node.nodeId, "worker", "bzzz", "announce", "/capabilities", "",
		)
		
		capabilities := []byte(fmt.Sprintf(`{
			"node_id": "%s",
			"address": "192.168.1.%d:8080",
			"vram_gb": %d,
			"available_models": %v,
			"status": "online",
			"load": 0.1
		}`, node.nodeId, 110+i, node.vram, fmt.Sprintf("%q", node.models)))
		
		err := mockDHT.PutValue(ctx, nodeAddr, capabilities)
		if err != nil {
			t.Fatalf("Failed to store node capabilities for %s: %v", node.nodeId, err)
		}
		
		// Announce as provider for model processing
		err = mockDHT.Provide(ctx, "bzzz:models:processing", node.nodeId)
		if err != nil {
			t.Fatalf("Failed to announce model provider %s: %v", node.nodeId, err)
		}
	}
	
	// Step 2: RUSTLE browser queries for available models
	modelQueryAddr, _ := ucxl.GenerateUCXLAddress(
		"rustle-browser", "client", "bzzz", "query", "/models/list", "^/",
	)
	
	// Step 3: Find all model processing providers
	modelProviders, err := mockDHT.FindProviders(ctx, "bzzz:models:processing")
	if err != nil {
		t.Fatalf("Failed to find model providers: %v", err)
	}
	
	if len(modelProviders) != len(processingNodes) {
		t.Fatalf("Expected %d model providers, got %d", len(processingNodes), len(modelProviders))
	}
	
	// Step 4: Aggregate model information (simulating coordinator behavior)
	aggregatedModels := []byte(`{
		"available_models": [
			{"name": "llama-3.1-8b", "node": "worker-001", "vram_required": 8},
			{"name": "qwen-2.5-7b", "node": "worker-001", "vram_required": 7},
			{"name": "llama-3.1-70b", "node": "worker-002", "vram_required": 70},
			{"name": "mixtral-8x7b", "node": "worker-003", "vram_required": 32},
			{"name": "qwen-2.5-32b", "node": "worker-003", "vram_required": 28}
		],
		"total_nodes": 3,
		"total_models": 5,
		"query_timestamp": "` + time.Now().Format(time.RFC3339) + `"
	}`)
	
	err = mockDHT.PutValue(ctx, modelQueryAddr, aggregatedModels)
	if err != nil {
		t.Fatalf("Failed to store aggregated model information: %v", err)
	}
	
	// Step 5: RUSTLE retrieves the aggregated information
	retrieved, err := mockDHT.GetValue(ctx, modelQueryAddr)
	if err != nil {
		t.Fatalf("Failed to retrieve model information: %v", err)
	}
	
	if len(retrieved) == 0 {
		t.Fatalf("Retrieved model information is empty")
	}
	
	t.Logf("✓ RUSTLE Model Discovery Scenario: Successfully discovered models")
	t.Logf("  - Processing nodes: %d", len(processingNodes))
	t.Logf("  - Model providers: %d", len(modelProviders))
	t.Logf("  - Model info size: %d bytes", len(retrieved))
}

// Benchmark tests for performance validation
func BenchmarkMockDHTOperations(b *testing.B) {
	mockDHT := dht.NewMockDHT()
	ctx := context.Background()
	
	b.Run("PutValue", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("benchmark-key-%d", i)
			value := []byte("benchmark-value")
			mockDHT.PutValue(ctx, key, value)
		}
	})
	
	b.Run("GetValue", func(b *testing.B) {
		// Pre-populate
		for i := 0; i < 1000; i++ {
			key := fmt.Sprintf("benchmark-key-%d", i)
			value := []byte("benchmark-value")
			mockDHT.PutValue(ctx, key, value)
		}
		
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("benchmark-key-%d", i%1000)
			mockDHT.GetValue(ctx, key)
		}
	})
}

func BenchmarkUCXLAddressOperations(b *testing.B) {
	b.Run("ParseAddress", func(b *testing.B) {
		address := "ucxl://agent:role@project:task/path*temporal/"
		for i := 0; i < b.N; i++ {
			ucxl.ParseUCXLAddress(address)
		}
	})
	
	b.Run("GenerateAddress", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			ucxl.GenerateUCXLAddress("agent", "role", "project", "task", "/path", "temporal")
		}
	})
}