package server

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

	"github.com/CalmoAI/mcp-server-signoz/internal/domain"
	"github.com/CalmoAI/mcp-server-signoz/pkg/mcp"
)

// Test Transport for testing
type TestTransport struct {
	sent    []*mcp.MCPMessage
	toSend  []*mcp.MCPMessage
	sendIdx int
	closed  bool
}

func NewTestTransport() *TestTransport {
	return &TestTransport{
		sent:   make([]*mcp.MCPMessage, 0),
		toSend: make([]*mcp.MCPMessage, 0),
	}
}

func (t *TestTransport) Send(message *mcp.MCPMessage) error {
	if t.closed {
		return io.EOF
	}
	t.sent = append(t.sent, message)
	return nil
}

func (t *TestTransport) Receive() (*mcp.MCPMessage, error) {
	if t.closed {
		return nil, io.EOF
	}
	if t.sendIdx >= len(t.toSend) {
		return nil, io.EOF
	}
	msg := t.toSend[t.sendIdx]
	t.sendIdx++
	return msg, nil
}

func (t *TestTransport) Close() error {
	t.closed = true
	return nil
}

func (t *TestTransport) QueueMessage(msg *mcp.MCPMessage) {
	t.toSend = append(t.toSend, msg)
}

func (t *TestTransport) GetSentMessages() []*mcp.MCPMessage {
	return t.sent
}

// Mock implementations for testing
type MockMCPService struct{}

func (m *MockMCPService) TestConnection(ctx context.Context) (*domain.ConnectionStatus, error) {
	return &domain.ConnectionStatus{
		Connected:   true,
		URL:         "http://localhost:3301",
		Version:     "1.0.0",
		LastChecked: time.Now(),
		Latency:     time.Millisecond * 50,
	}, nil
}

func (m *MockMCPService) FetchDashboards(ctx context.Context) ([]domain.Dashboard, error) {
	return []domain.Dashboard{
		{
			ID:          "dashboard-1",
			Name:        "Test Dashboard",
			Description: "A test dashboard",
			Tags:        []string{"test", "example"},
			CreatedAt:   time.Now().Add(-24 * time.Hour),
			UpdatedAt:   time.Now(),
		},
	}, nil
}

func (m *MockMCPService) FetchDashboardDetails(ctx context.Context, dashboardID string) (*domain.Dashboard, error) {
	return &domain.Dashboard{
		ID:          dashboardID,
		Name:        "Test Dashboard Details",
		Description: "Detailed test dashboard",
		Tags:        []string{"test", "details"},
		CreatedAt:   time.Now().Add(-24 * time.Hour),
		UpdatedAt:   time.Now(),
	}, nil
}

func (m *MockMCPService) FetchDashboardData(ctx context.Context, dashboardName string, startTime, endTime time.Time) (*domain.DashboardData, error) {
	return &domain.DashboardData{
		DashboardID: "dashboard-1",
		TimeRange: domain.TimeRange{
			Start: startTime,
			End:   endTime,
		},
		Timestamp: time.Now(),
		Panels: []domain.PanelData{
			{
				PanelID: "panel-1",
				Data:    map[string]interface{}{"value": 42},
				Status:  "success",
			},
		},
	}, nil
}

func (m *MockMCPService) FetchAPMMetrics(ctx context.Context, serviceName string, startTime, endTime time.Time) (*domain.APMMetrics, error) {
	return &domain.APMMetrics{
		ServiceName: serviceName,
		TimeRange: domain.TimeRange{
			Start: startTime,
			End:   endTime,
		},
		Metrics: map[string]domain.Metric{
			"requests_per_minute": {
				Name:        "requests_per_minute",
				Unit:        "requests/min",
				Description: "Number of requests per minute",
				DataPoints: []domain.DataPoint{
					{Timestamp: startTime, Value: 100},
					{Timestamp: endTime, Value: 110},
				},
			},
		},
	}, nil
}

func (m *MockMCPService) FetchServices(ctx context.Context, startTime, endTime *time.Time) ([]domain.Service, error) {
	return []domain.Service{
		{
			ServiceName: "test-service",
			Environment: "test",
			Tags:        map[string]string{"version": "1.0.0"},
			Operations:  []string{"GET /api/users", "POST /api/orders"},
			CallRate:    100.5,
			ErrorRate:   0.02,
			P99Latency:  250.3,
			AvgLatency:  125.7,
			LastSeen:    time.Now(),
		},
	}, nil
}

func (m *MockMCPService) ExecuteClickHouseQuery(ctx context.Context, sql string, startTime, endTime time.Time, params map[string]interface{}) (*domain.QueryResult, error) {
	return &domain.QueryResult{
		Columns: []domain.ColumnInfo{
			{Name: "timestamp", Type: "UInt64"},
			{Name: "value", Type: "Float64"},
		},
		Rows: [][]interface{}{
			{startTime.Unix(), 42},
			{endTime.Unix(), 43},
		},
		Meta: domain.QueryMeta{
			QueryID:       "query-123",
			ExecutionTime: time.Millisecond * 25,
			RowsAffected:  2,
			Statistics: domain.QueryStats{
				ReadRows:  2,
				ReadBytes: 64,
			},
		},
	}, nil
}

func (m *MockMCPService) ExecuteBuilderQuery(ctx context.Context, queryType, metricName, aggregation string, filters []domain.QueryFilter, groupBy []string, startTime, endTime time.Time, step time.Duration, options map[string]interface{}) (*domain.QueryResult, error) {
	return &domain.QueryResult{
		Columns: []domain.ColumnInfo{
			{Name: "timestamp", Type: "UInt64"},
			{Name: "value", Type: "Float64"},
		},
		Rows: [][]interface{}{
			{startTime.Unix(), 100},
			{endTime.Unix(), 110},
		},
		Meta: domain.QueryMeta{
			QueryID:       "builder-query-456",
			ExecutionTime: time.Millisecond * 18,
			RowsAffected:  2,
			Statistics: domain.QueryStats{
				ReadRows:  2,
				ReadBytes: 48,
			},
		},
	}, nil
}

func (m *MockMCPService) FetchTracesOrLogs(ctx context.Context, dataType string, startTime, endTime time.Time, serviceName string, limit int) ([]domain.TraceLog, error) {
	return []domain.TraceLog{
		{
			Type:      dataType,
			ID:        "trace-1",
			Timestamp: startTime,
			Service:   serviceName,
			Operation: "test_operation",
			Duration:  time.Millisecond * 125,
			Status:    "success",
			Message:   "Test trace message",
			Attributes: map[string]interface{}{
				"operation": "test_operation",
			},
			SpanID:   "span-123",
			TraceID:  "trace-abc",
			ParentID: "",
		},
	}, nil
}

type MockLogger struct{}

func (m *MockLogger) Debug(msg string, fields ...domain.Field)  {}
func (m *MockLogger) Info(msg string, fields ...domain.Field)   {}
func (m *MockLogger) Warn(msg string, fields ...domain.Field)   {}
func (m *MockLogger) Error(msg string, fields ...domain.Field)  {}
func (m *MockLogger) Fatal(msg string, fields ...domain.Field)  {}
func (m *MockLogger) With(fields ...domain.Field) domain.Logger { return m }

type MockTracer struct{}

func (m *MockTracer) StartSpan(ctx context.Context, name string, opts ...domain.SpanOption) (context.Context, domain.Span) {
	return ctx, &MockSpan{}
}

type MockSpan struct{}

func (m *MockSpan) SetAttribute(key string, value interface{})        {}
func (m *MockSpan) SetStatus(code domain.StatusCode, message string)  {}
func (m *MockSpan) RecordError(err error, opts ...domain.EventOption) {}
func (m *MockSpan) End()                                              {}

type MockMetrics struct{}

func (m *MockMetrics) Counter(name string) domain.Counter {
	return &MockCounter{}
}

func (m *MockMetrics) Histogram(name string) domain.Histogram {
	return &MockHistogram{}
}

func (m *MockMetrics) Gauge(name string) domain.Gauge {
	return &MockGauge{}
}

type MockCounter struct{}

func (m *MockCounter) Add(ctx context.Context, value int64, attributes ...domain.Attribute) {}

type MockHistogram struct{}

func (m *MockHistogram) Record(ctx context.Context, value float64, attributes ...domain.Attribute) {}

type MockGauge struct{}

func (m *MockGauge) Set(ctx context.Context, value float64, attributes ...domain.Attribute) {}

// Helper functions for tests
func createTestServer() *Server {
	transport := NewTestTransport()
	mcpService := &MockMCPService{}
	logger := &MockLogger{}
	tracer := &MockTracer{}
	metrics := &MockMetrics{}

	return NewServer(transport, mcpService, logger, tracer, metrics)
}

func createInitializeMessage() *mcp.MCPMessage {
	initRequest := mcp.InitializeRequest{
		ProtocolVersion: mcp.MCPVersion,
		Capabilities:    mcp.Capabilities{},
		ClientInfo: mcp.ClientInfo{
			Name:    "test-client",
			Version: "1.0.0",
		},
	}

	params, _ := json.Marshal(initRequest)
	id := int64(1)

	return &mcp.MCPMessage{
		JSONRPC: "2.0",
		ID:      &id,
		Method:  "initialize",
		Params:  params,
	}
}

func createToolsListMessage() *mcp.MCPMessage {
	id := int64(2)
	return &mcp.MCPMessage{
		JSONRPC: "2.0",
		ID:      &id,
		Method:  "tools/list",
	}
}

func createToolsCallMessage(toolName string, args map[string]interface{}) *mcp.MCPMessage {
	toolsCallRequest := mcp.ToolsCallRequest{
		Name:      toolName,
		Arguments: args,
	}

	params, _ := json.Marshal(toolsCallRequest)
	id := int64(3)

	return &mcp.MCPMessage{
		JSONRPC: "2.0",
		ID:      &id,
		Method:  "tools/call",
		Params:  params,
	}
}

// Test MCP Protocol Compliance
func TestMCPProtocolCompliance(t *testing.T) {
	server := createTestServer()
	transport := server.transport.(*TestTransport)

	// Test Initialize
	t.Run("Initialize", func(t *testing.T) {
		initMsg := createInitializeMessage()
		transport.QueueMessage(initMsg)

		// Process the message
		server.handleMessage(initMsg)

		// Check response
		sentMessages := transport.GetSentMessages()
		if len(sentMessages) != 1 {
			t.Errorf("Expected 1 sent message, got %d", len(sentMessages))
			return
		}

		response := sentMessages[0]
		if response.JSONRPC != "2.0" {
			t.Errorf("Expected JSONRPC 2.0, got %s", response.JSONRPC)
		}

		if response.ID == nil || *response.ID != 1 {
			t.Errorf("Expected response ID 1, got %v", response.ID)
		}

		if response.Error != nil {
			t.Errorf("Expected no error, got %v", response.Error)
		}

		// Verify response contains server info and capabilities
		var initResponse mcp.InitializeResponse
		if err := json.Unmarshal(response.Result, &initResponse); err != nil {
			t.Errorf("Failed to unmarshal initialize response: %v", err)
		}

		if initResponse.ServerInfo.Name != "signoz-mcp-server" {
			t.Errorf("Expected server name 'signoz-mcp-server', got %s", initResponse.ServerInfo.Name)
		}
	})

	// Test Tools List
	t.Run("ToolsList", func(t *testing.T) {
		// Reset transport
		transport.sent = nil
		toolsListMsg := createToolsListMessage()

		server.handleMessage(toolsListMsg)

		sentMessages := transport.GetSentMessages()
		if len(sentMessages) != 1 {
			t.Errorf("Expected 1 sent message, got %d", len(sentMessages))
			return
		}

		response := sentMessages[0]
		if response.Error != nil {
			t.Errorf("Expected no error, got %v", response.Error)
		}

		// Verify tools list response
		var toolsResponse mcp.ToolsListResponse
		if err := json.Unmarshal(response.Result, &toolsResponse); err != nil {
			t.Errorf("Failed to unmarshal tools response: %v", err)
		}

		expectedToolCount := 9
		if len(toolsResponse.Tools) != expectedToolCount {
			t.Errorf("Expected %d tools, got %d", expectedToolCount, len(toolsResponse.Tools))
		}

		// Verify specific tools exist
		expectedTools := []string{
			"signoz_test_connection",
			"signoz_fetch_dashboards",
			"signoz_fetch_dashboard_details",
			"signoz_fetch_dashboard_data",
			"signoz_fetch_apm_metrics",
			"signoz_fetch_services",
			"signoz_execute_clickhouse_query",
			"signoz_execute_builder_query",
			"signoz_fetch_traces_or_logs",
		}

		toolNames := make(map[string]bool)
		for _, tool := range toolsResponse.Tools {
			toolNames[tool.Name] = true
		}

		for _, expectedTool := range expectedTools {
			if !toolNames[expectedTool] {
				t.Errorf("Expected tool %s not found in tools list", expectedTool)
			}
		}
	})
}

// Test Individual Tools
func TestToolExecution(t *testing.T) {
	server := createTestServer()
	transport := server.transport.(*TestTransport)

	// Initialize server first
	initMsg := createInitializeMessage()
	server.handleMessage(initMsg)
	transport.sent = nil // Reset sent messages

	testCases := []struct {
		name     string
		toolName string
		args     map[string]interface{}
		wantErr  bool
	}{
		{
			name:     "TestConnection",
			toolName: "signoz_test_connection",
			args:     map[string]interface{}{},
			wantErr:  false,
		},
		{
			name:     "FetchDashboards",
			toolName: "signoz_fetch_dashboards",
			args:     map[string]interface{}{},
			wantErr:  false,
		},
		{
			name:     "FetchDashboardDetails",
			toolName: "signoz_fetch_dashboard_details",
			args: map[string]interface{}{
				"dashboard_id": "test-dashboard-id",
			},
			wantErr: false,
		},
		{
			name:     "FetchDashboardData",
			toolName: "signoz_fetch_dashboard_data",
			args: map[string]interface{}{
				"dashboard_name": "test-dashboard",
				"start_time":     "2023-01-01T00:00:00Z",
				"end_time":       "2023-01-01T01:00:00Z",
			},
			wantErr: false,
		},
		{
			name:     "FetchAPMMetrics",
			toolName: "signoz_fetch_apm_metrics",
			args: map[string]interface{}{
				"service_name": "test-service",
				"start_time":   "2023-01-01T00:00:00Z",
				"end_time":     "2023-01-01T01:00:00Z",
			},
			wantErr: false,
		},
		{
			name:     "FetchServices",
			toolName: "signoz_fetch_services",
			args:     map[string]interface{}{},
			wantErr:  false,
		},
		{
			name:     "ExecuteClickHouseQuery",
			toolName: "signoz_execute_clickhouse_query",
			args: map[string]interface{}{
				"sql":        "SELECT * FROM traces LIMIT 10",
				"start_time": "2023-01-01T00:00:00Z",
				"end_time":   "2023-01-01T01:00:00Z",
			},
			wantErr: false,
		},
		{
			name:     "ExecuteBuilderQuery",
			toolName: "signoz_execute_builder_query",
			args: map[string]interface{}{
				"query_type":  "metrics",
				"metric_name": "http_requests_total",
				"aggregation": "sum",
				"start_time":  "2023-01-01T00:00:00Z",
				"end_time":    "2023-01-01T01:00:00Z",
			},
			wantErr: false,
		},
		{
			name:     "FetchTracesOrLogs",
			toolName: "signoz_fetch_traces_or_logs",
			args: map[string]interface{}{
				"data_type":    "traces",
				"service_name": "test-service",
				"start_time":   "2023-01-01T00:00:00Z",
				"end_time":     "2023-01-01T01:00:00Z",
				"limit":        float64(100),
			},
			wantErr: false,
		},
		{
			name:     "UnknownTool",
			toolName: "unknown_tool",
			args:     map[string]interface{}{},
			wantErr:  true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			transport.sent = nil // Reset for each test

			toolCallMsg := createToolsCallMessage(tc.toolName, tc.args)
			server.handleMessage(toolCallMsg)

			sentMessages := transport.GetSentMessages()
			if len(sentMessages) != 1 {
				t.Errorf("Expected 1 sent message, got %d", len(sentMessages))
				return
			}

			response := sentMessages[0]

			if tc.wantErr {
				// For error cases, check that we get an error response
				var toolResponse mcp.ToolsCallResponse
				if err := json.Unmarshal(response.Result, &toolResponse); err != nil {
					t.Errorf("Failed to unmarshal tool response: %v", err)
					return
				}

				if !toolResponse.IsError {
					t.Errorf("Expected error response for tool %s", tc.toolName)
				}
			} else {
				// For success cases, check that we get a valid response
				if response.Error != nil {
					t.Errorf("Unexpected error for tool %s: %v", tc.toolName, response.Error)
					return
				}

				var toolResponse mcp.ToolsCallResponse
				if err := json.Unmarshal(response.Result, &toolResponse); err != nil {
					t.Errorf("Failed to unmarshal tool response: %v", err)
					return
				}

				if toolResponse.IsError {
					t.Errorf("Unexpected error response for tool %s", tc.toolName)
				}

				if len(toolResponse.Content) == 0 {
					t.Errorf("Expected content in response for tool %s", tc.toolName)
				}
			}
		})
	}
}

// Test Error Handling
func TestErrorHandling(t *testing.T) {
	server := createTestServer()
	transport := server.transport.(*TestTransport)

	t.Run("InvalidJSON", func(t *testing.T) {
		id := int64(1)
		invalidMsg := &mcp.MCPMessage{
			JSONRPC: "2.0",
			ID:      &id,
			Method:  "tools/call",
			Params:  json.RawMessage(`{"invalid": json`), // Invalid JSON
		}

		server.handleMessage(invalidMsg)

		sentMessages := transport.GetSentMessages()
		if len(sentMessages) != 1 {
			t.Errorf("Expected 1 sent message, got %d", len(sentMessages))
			return
		}

		response := sentMessages[0]
		if response.Error == nil {
			t.Errorf("Expected error response for invalid JSON")
		}
	})

	t.Run("MissingRequiredField", func(t *testing.T) {
		transport.sent = nil

		// Initialize server first
		initMsg := createInitializeMessage()
		server.handleMessage(initMsg)
		transport.sent = nil

		// Call tool with missing required field
		toolCallMsg := createToolsCallMessage("signoz_fetch_dashboard_details", map[string]interface{}{
			// Missing dashboard_id
		})

		server.handleMessage(toolCallMsg)

		sentMessages := transport.GetSentMessages()
		if len(sentMessages) != 1 {
			t.Errorf("Expected 1 sent message, got %d", len(sentMessages))
			return
		}

		response := sentMessages[0]

		var toolResponse mcp.ToolsCallResponse
		if err := json.Unmarshal(response.Result, &toolResponse); err != nil {
			t.Errorf("Failed to unmarshal tool response: %v", err)
			return
		}

		if !toolResponse.IsError {
			t.Errorf("Expected error response for missing required field")
		}
	})
}

// Test Tool Input Schema Validation
func TestToolInputSchemas(t *testing.T) {
	server := createTestServer()

	// Get tools from server
	expectedSchemas := map[string]struct {
		hasRequiredFields bool
		requiredFields    []string
	}{
		"signoz_test_connection":          {false, []string{}},
		"signoz_fetch_dashboards":         {false, []string{}},
		"signoz_fetch_dashboard_details":  {true, []string{"dashboard_id"}},
		"signoz_fetch_dashboard_data":     {true, []string{"dashboard_name", "start_time", "end_time"}},
		"signoz_fetch_apm_metrics":        {true, []string{"service_name", "start_time", "end_time"}},
		"signoz_fetch_services":           {false, []string{}},
		"signoz_execute_clickhouse_query": {true, []string{"sql", "start_time", "end_time"}},
		"signoz_execute_builder_query":    {true, []string{"query_type", "metric_name", "aggregation", "start_time", "end_time"}},
		"signoz_fetch_traces_or_logs":     {true, []string{"data_type", "service_name", "start_time", "end_time", "limit"}},
	}

	for toolName, expectedSchema := range expectedSchemas {
		t.Run(toolName, func(t *testing.T) {
			tool, exists := server.tools[toolName]
			if !exists {
				t.Errorf("Tool %s not found in server tools", toolName)
				return
			}

			if expectedSchema.hasRequiredFields {
				if len(tool.InputSchema.Required) != len(expectedSchema.requiredFields) {
					t.Errorf("Tool %s: expected %d required fields, got %d",
						toolName, len(expectedSchema.requiredFields), len(tool.InputSchema.Required))
				}

				// Check all required fields are present
				requiredMap := make(map[string]bool)
				for _, field := range tool.InputSchema.Required {
					requiredMap[field] = true
				}

				for _, expectedField := range expectedSchema.requiredFields {
					if !requiredMap[expectedField] {
						t.Errorf("Tool %s: missing required field %s", toolName, expectedField)
					}
				}
			} else {
				if len(tool.InputSchema.Required) != 0 {
					t.Errorf("Tool %s: expected no required fields, got %d",
						toolName, len(tool.InputSchema.Required))
				}
			}

			// Verify input schema has proper structure
			if tool.InputSchema.Type != "object" {
				t.Errorf("Tool %s: expected input schema type 'object', got %s",
					toolName, tool.InputSchema.Type)
			}
		})
	}
}
