package mcp

import (
	"encoding/json"
	"strconv"
	"testing"
	"time"
)

// Benchmark MCP message serialization/deserialization
func BenchmarkMessageSerialization(b *testing.B) {
	message := &MCPMessage{
		JSONRPC: "2.0",
		ID:      func() *int64 { v := int64(1); return &v }(),
		Method:  "tools/call",
		Params:  json.RawMessage(`{"name":"signoz_test_connection","arguments":{}}`),
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(message)
		if err != nil {
			b.Fatal(err)
		}

		var parsed MCPMessage
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkInitializeRequest(b *testing.B) {
	initRequest := InitializeRequest{
		ProtocolVersion: MCPVersion,
		Capabilities:    Capabilities{},
		ClientInfo: ClientInfo{
			Name:    "benchmark-client",
			Version: "1.0.0",
		},
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(initRequest)
		if err != nil {
			b.Fatal(err)
		}

		var parsed InitializeRequest
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkToolsListResponse(b *testing.B) {
	tools := make([]Tool, 10)
	for i := 0; i < 10; i++ {
		tools[i] = Tool{
			Name:        "test_tool_" + strconv.Itoa(i),
			Description: "A test tool for benchmarking purposes",
			InputSchema: ToolInputSchema{
				Type: "object",
				Properties: map[string]Property{
					"param1": {
						Type:        "string",
						Description: "First parameter",
					},
					"param2": {
						Type:        "number",
						Description: "Second parameter",
					},
				},
				Required: []string{"param1"},
			},
		}
	}

	response := ToolsListResponse{
		Tools: tools,
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(response)
		if err != nil {
			b.Fatal(err)
		}

		var parsed ToolsListResponse
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkToolsCallRequest(b *testing.B) {
	request := ToolsCallRequest{
		Name: "signoz_execute_clickhouse_query",
		Arguments: map[string]interface{}{
			"sql":        "SELECT * FROM traces WHERE service_name = ? AND timestamp >= ? AND timestamp <= ? LIMIT ?",
			"start_time": time.Now().Add(-time.Hour).Format(time.RFC3339),
			"end_time":   time.Now().Format(time.RFC3339),
			"params": map[string]interface{}{
				"service_name": "user-service",
				"limit":        1000,
			},
		},
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(request)
		if err != nil {
			b.Fatal(err)
		}

		var parsed ToolsCallRequest
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkToolsCallResponse(b *testing.B) {
	// Simulate a large response with trace data
	content := []Content{
		{
			Type: string(ContentTypeText),
			Text: `{
  "columns": ["timestamp", "trace_id", "span_id", "service_name", "operation", "duration_ms", "status_code"],
  "rows": [
    [1642694400, "abc123", "def456", "user-service", "GET /api/users", 125.5, 200],
    [1642694401, "abc124", "def457", "user-service", "GET /api/users/123", 89.2, 200],
    [1642694402, "abc125", "def458", "order-service", "POST /api/orders", 234.7, 201],
    [1642694403, "abc126", "def459", "payment-service", "POST /api/payments", 456.3, 200],
    [1642694404, "abc127", "def460", "user-service", "PUT /api/users/123", 178.9, 200]
  ],
  "row_count": 5,
  "execution_ms": 42,
  "bytes_read": 2048
}`,
		},
	}

	response := ToolsCallResponse{
		Content: content,
		IsError: false,
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(response)
		if err != nil {
			b.Fatal(err)
		}

		var parsed ToolsCallResponse
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkErrorResponse(b *testing.B) {
	errorResp := MCPError{
		Code:    ErrorCodeInternalError,
		Message: "Internal server error occurred while processing the request",
		Data:    json.RawMessage(`{"details":"Database connection failed","error_code":"DB_CONNECTION_ERROR"}`),
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(errorResp)
		if err != nil {
			b.Fatal(err)
		}

		var parsed MCPError
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// Benchmark concurrent message processing
func BenchmarkConcurrentMessageProcessing(b *testing.B) {
	messages := make([]*MCPMessage, 100)
	for i := 0; i < 100; i++ {
		id := int64(i)
		messages[i] = &MCPMessage{
			JSONRPC: "2.0",
			ID:      &id,
			Method:  "tools/call",
			Params:  json.RawMessage(`{"name":"signoz_test_connection","arguments":{}}`),
		}
	}

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		i := 0
		for pb.Next() {
			message := messages[i%len(messages)]
			data, err := json.Marshal(message)
			if err != nil {
				b.Fatal(err)
			}

			var parsed MCPMessage
			err = json.Unmarshal(data, &parsed)
			if err != nil {
				b.Fatal(err)
			}
			i++
		}
	})
}

// Benchmark message validation
func BenchmarkMessageValidation(b *testing.B) {
	validMessage := &MCPMessage{
		JSONRPC: "2.0",
		ID:      func() *int64 { v := int64(1); return &v }(),
		Method:  "initialize",
		Params:  json.RawMessage(`{"protocolVersion":"2025-06-18"}`),
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		// Simulate validation logic
		if validMessage.JSONRPC != "2.0" {
			b.Fatal("Invalid JSONRPC version")
		}
		if validMessage.ID == nil {
			b.Fatal("Missing ID")
		}
		if validMessage.Method == "" {
			b.Fatal("Missing method")
		}
	}
}

// Benchmark large payload handling
func BenchmarkLargePayload(b *testing.B) {
	// Create a large payload simulating a dashboard with many panels
	largePayload := map[string]interface{}{
		"dashboard_name": "large-dashboard",
		"panels":         make([]map[string]interface{}, 100),
	}

	for i := 0; i < 100; i++ {
		panel := map[string]interface{}{
			"panel_id":    "panel-" + strconv.Itoa(i),
			"title":       "Panel " + strconv.Itoa(i),
			"description": "A detailed panel with lots of data points and configuration options",
			"data": map[string]interface{}{
				"series": make([]map[string]interface{}, 10),
			},
		}

		series := panel["data"].(map[string]interface{})["series"].([]map[string]interface{})
		for j := 0; j < 10; j++ {
			series[j] = map[string]interface{}{
				"name":   "series-" + strconv.Itoa(j),
				"values": make([][]interface{}, 1000),
			}
			values := series[j]["values"].([][]interface{})
			for k := 0; k < 1000; k++ {
				values[k] = []interface{}{
					time.Now().Add(time.Duration(k) * time.Minute).Unix(),
					float64(k) * 1.5,
				}
			}
		}

		largePayload["panels"].([]map[string]interface{})[i] = panel
	}

	content := []Content{
		{
			Type: string(ContentTypeText),
		},
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(largePayload)
		if err != nil {
			b.Fatal(err)
		}

		content[0].Text = string(data)

		response := ToolsCallResponse{
			Content: content,
			IsError: false,
		}

		responseData, err := json.Marshal(response)
		if err != nil {
			b.Fatal(err)
		}

		var parsed ToolsCallResponse
		err = json.Unmarshal(responseData, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// Memory allocation benchmarks
func BenchmarkMemoryAllocation(b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		message := &MCPMessage{
			JSONRPC: "2.0",
			ID:      func() *int64 { v := int64(i); return &v }(),
			Method:  "tools/call",
			Params:  json.RawMessage(`{"name":"test","arguments":{}}`),
		}

		_ = message // Use the message to prevent optimization
	}
}

func BenchmarkJSONMarshaling(b *testing.B) {
	testData := map[string]interface{}{
		"string_field":  "test value",
		"number_field":  42.5,
		"boolean_field": true,
		"array_field":   []interface{}{1, 2, 3, "four", 5.0},
		"object_field": map[string]interface{}{
			"nested_string": "nested value",
			"nested_number": 123,
		},
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data, err := json.Marshal(testData)
		if err != nil {
			b.Fatal(err)
		}

		var parsed map[string]interface{}
		err = json.Unmarshal(data, &parsed)
		if err != nil {
			b.Fatal(err)
		}
	}
}
