package transport

import (
	"bytes"
	"encoding/json"
	"io"
	"strings"
	"testing"
	"time"

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

func TestStdioTransport(t *testing.T) {
	t.Run("NewStdioTransport", func(t *testing.T) {
		transport := NewStdioTransport()
		if transport == nil {
			t.Error("Expected non-nil transport")
		}

		if transport.IsClosed() {
			t.Error("Expected transport to not be closed initially")
		}
	})

	t.Run("NewStdioTransportWithStreams", func(t *testing.T) {
		stdin := strings.NewReader("")
		stdout := &bytes.Buffer{}

		transport := NewStdioTransportWithStreams(stdin, stdout)
		if transport == nil {
			t.Error("Expected non-nil transport")
		}

		if transport.stdin != stdin {
			t.Error("Expected stdin to be set correctly")
		}

		if transport.stdout != stdout {
			t.Error("Expected stdout to be set correctly")
		}
	})
}

func TestStdioTransportSend(t *testing.T) {
	testCases := []struct {
		name    string
		message *mcp.MCPMessage
		wantErr bool
	}{
		{
			name: "ValidRequest",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Method:  "initialize",
				Params:  json.RawMessage(`{"protocolVersion":"2025-06-18"}`),
			},
			wantErr: false,
		},
		{
			name: "ValidResponse",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Result:  json.RawMessage(`{"success":true}`),
			},
			wantErr: false,
		},
		{
			name: "ValidNotification",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				Method:  "notification",
				Params:  json.RawMessage(`{"type":"update"}`),
			},
			wantErr: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			stdout := &bytes.Buffer{}
			transport := NewStdioTransportWithStreams(nil, stdout)

			err := transport.Send(tc.message)

			if tc.wantErr && err == nil {
				t.Error("Expected error but got none")
			}

			if !tc.wantErr && err != nil {
				t.Errorf("Expected no error but got: %v", err)
			}

			if !tc.wantErr {
				// Verify the message was written correctly
				output := stdout.String()

				// Should end with newline
				if !strings.HasSuffix(output, "\n") {
					t.Error("Expected output to end with newline")
				}

				// Should be valid JSON
				var parsed mcp.MCPMessage
				if err := json.Unmarshal([]byte(strings.TrimSpace(output)), &parsed); err != nil {
					t.Errorf("Output is not valid JSON: %v", err)
				}

				// Should match original message
				if parsed.JSONRPC != tc.message.JSONRPC {
					t.Errorf("Expected JSONRPC %s, got %s", tc.message.JSONRPC, parsed.JSONRPC)
				}

				if tc.message.ID != nil {
					if parsed.ID == nil || *parsed.ID != *tc.message.ID {
						t.Errorf("ID mismatch: expected %v, got %v", tc.message.ID, parsed.ID)
					}
				}

				if parsed.Method != tc.message.Method {
					t.Errorf("Expected method %s, got %s", tc.message.Method, parsed.Method)
				}
			}
		})
	}
}

func TestStdioTransportReceive(t *testing.T) {
	testCases := []struct {
		name     string
		input    string
		expected *mcp.MCPMessage
		wantErr  bool
	}{
		{
			name:  "ValidRequest",
			input: `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}`,
			expected: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Method:  "initialize",
				Params:  json.RawMessage(`{"protocolVersion":"2025-06-18"}`),
			},
			wantErr: false,
		},
		{
			name:  "ValidResponse",
			input: `{"jsonrpc":"2.0","id":1,"result":{"success":true}}`,
			expected: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Result:  json.RawMessage(`{"success":true}`),
			},
			wantErr: false,
		},
		{
			name:  "ValidNotification",
			input: `{"jsonrpc":"2.0","method":"notification","params":{"type":"update"}}`,
			expected: &mcp.MCPMessage{
				JSONRPC: "2.0",
				Method:  "notification",
				Params:  json.RawMessage(`{"type":"update"}`),
			},
			wantErr: false,
		},
		{
			name:    "InvalidJSON",
			input:   `{"jsonrpc":"2.0","id":1,"method":}`,
			wantErr: true,
		},
		{
			name:    "InvalidJSONRPCVersion",
			input:   `{"jsonrpc":"1.0","id":1,"method":"test"}`,
			wantErr: true,
		},
		{
			name:  "EmptyLine",
			input: "",
			// Should skip empty lines and return EOF
			wantErr: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			stdin := strings.NewReader(tc.input + "\n")
			transport := NewStdioTransportWithStreams(stdin, nil)

			message, err := transport.Receive()

			if tc.wantErr {
				if err == nil {
					t.Error("Expected error but got none")
				}
				return
			}

			if err != nil {
				t.Errorf("Expected no error but got: %v", err)
				return
			}

			if message == nil {
				t.Error("Expected non-nil message")
				return
			}

			// Verify message fields
			if message.JSONRPC != tc.expected.JSONRPC {
				t.Errorf("Expected JSONRPC %s, got %s", tc.expected.JSONRPC, message.JSONRPC)
			}

			if tc.expected.ID != nil {
				if message.ID == nil || *message.ID != *tc.expected.ID {
					t.Errorf("ID mismatch: expected %v, got %v", tc.expected.ID, message.ID)
				}
			} else if message.ID != nil {
				t.Errorf("Expected nil ID, got %v", message.ID)
			}

			if message.Method != tc.expected.Method {
				t.Errorf("Expected method %s, got %s", tc.expected.Method, message.Method)
			}

			// Compare JSON content for params and result
			if tc.expected.Params != nil {
				if message.Params == nil {
					t.Error("Expected non-nil params")
				} else {
					var expectedParams, actualParams map[string]interface{}
					json.Unmarshal(tc.expected.Params, &expectedParams)
					json.Unmarshal(message.Params, &actualParams)

					if len(expectedParams) != len(actualParams) {
						t.Errorf("Params length mismatch: expected %d, got %d", len(expectedParams), len(actualParams))
					}
				}
			}

			if tc.expected.Result != nil {
				if message.Result == nil {
					t.Error("Expected non-nil result")
				} else {
					var expectedResult, actualResult map[string]interface{}
					json.Unmarshal(tc.expected.Result, &expectedResult)
					json.Unmarshal(message.Result, &actualResult)

					if len(expectedResult) != len(actualResult) {
						t.Errorf("Result length mismatch: expected %d, got %d", len(expectedResult), len(actualResult))
					}
				}
			}
		})
	}
}

func TestStdioTransportEOF(t *testing.T) {
	stdin := strings.NewReader("")
	transport := NewStdioTransportWithStreams(stdin, nil)

	// Use a channel to handle timeout
	done := make(chan bool, 1)
	var err error

	go func() {
		_, err = transport.Receive()
		done <- true
	}()

	select {
	case <-done:
		if err != io.EOF {
			t.Errorf("Expected EOF error, got %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Errorf("Test timed out - transport.Receive() hung")
	}
}

func TestStdioTransportClose(t *testing.T) {
	stdout := &bytes.Buffer{}
	transport := NewStdioTransportWithStreams(nil, stdout)

	if transport.IsClosed() {
		t.Error("Expected transport to not be closed initially")
	}

	err := transport.Close()
	if err != nil {
		t.Errorf("Expected no error on close, got %v", err)
	}

	if !transport.IsClosed() {
		t.Error("Expected transport to be closed after Close()")
	}

	// Verify operations fail after close
	testMessage := &mcp.MCPMessage{
		JSONRPC: "2.0",
		Method:  "test",
	}

	err = transport.Send(testMessage)
	if err == nil {
		t.Error("Expected error when sending to closed transport")
	}

	_, err = transport.Receive()
	if err == nil {
		t.Error("Expected error when receiving from closed transport")
	}
}

func TestStdioTransportMultipleMessages(t *testing.T) {
	input := `{"jsonrpc":"2.0","id":1,"method":"initialize"}
{"jsonrpc":"2.0","id":2,"method":"tools/list"}
{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"test"}}
`
	stdin := strings.NewReader(input)
	transport := NewStdioTransportWithStreams(stdin, nil)

	expectedMethods := []string{"initialize", "tools/list", "tools/call"}
	expectedIDs := []int64{1, 2, 3}

	for i, expectedMethod := range expectedMethods {
		message, err := transport.Receive()
		if err != nil {
			t.Errorf("Message %d: unexpected error: %v", i+1, err)
			continue
		}

		if message.Method != expectedMethod {
			t.Errorf("Message %d: expected method %s, got %s", i+1, expectedMethod, message.Method)
		}

		if message.ID == nil || *message.ID != expectedIDs[i] {
			t.Errorf("Message %d: expected ID %d, got %v", i+1, expectedIDs[i], message.ID)
		}
	}

	// Should get EOF after all messages
	_, err := transport.Receive()
	if err != io.EOF {
		t.Errorf("Expected EOF after all messages, got %v", err)
	}
}

func TestMessageValidation(t *testing.T) {
	testCases := []struct {
		name    string
		message *mcp.MCPMessage
		wantErr bool
	}{
		{
			name: "ValidRequest",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Method:  "test",
			},
			wantErr: false,
		},
		{
			name: "ValidResponse",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Result:  json.RawMessage(`{}`),
			},
			wantErr: false,
		},
		{
			name: "ValidNotification",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				Method:  "notification",
			},
			wantErr: false,
		},
		{
			name:    "NilMessage",
			message: nil,
			wantErr: true,
		},
		{
			name: "InvalidJSONRPC",
			message: &mcp.MCPMessage{
				JSONRPC: "1.0",
				Method:  "test",
			},
			wantErr: true,
		},
		{
			name: "InvalidRequest",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				// Missing method
			},
			wantErr: true,
		},
		{
			name: "InvalidResponse",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				ID:      func() *int64 { v := int64(1); return &v }(),
				Method:  "test", // Responses shouldn't have methods
				Result:  json.RawMessage(`{}`),
			},
			wantErr: true,
		},
		{
			name: "InvalidNotification",
			message: &mcp.MCPMessage{
				JSONRPC: "2.0",
				// Missing method for notification
			},
			wantErr: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := ValidateMessage(tc.message)

			if tc.wantErr && err == nil {
				t.Error("Expected validation error but got none")
			}

			if !tc.wantErr && err != nil {
				t.Errorf("Expected no validation error but got: %v", err)
			}
		})
	}
}

func TestMessageBuilder(t *testing.T) {
	t.Run("BuildRequest", func(t *testing.T) {
		builder := NewMessageBuilder()
		params := map[string]interface{}{"test": "value"}

		message := builder.Request(1, "test").WithParams(params).Build()

		if message.JSONRPC != "2.0" {
			t.Errorf("Expected JSONRPC 2.0, got %s", message.JSONRPC)
		}

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

		if message.Method != "test" {
			t.Errorf("Expected method 'test', got %s", message.Method)
		}

		if message.Params == nil {
			t.Error("Expected non-nil params")
		}
	})

	t.Run("BuildResponse", func(t *testing.T) {
		builder := NewMessageBuilder()
		result := map[string]interface{}{"success": true}

		message := builder.Response(1).WithResult(result).Build()

		if message.JSONRPC != "2.0" {
			t.Errorf("Expected JSONRPC 2.0, got %s", message.JSONRPC)
		}

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

		if message.Result == nil {
			t.Error("Expected non-nil result")
		}
	})

	t.Run("BuildNotification", func(t *testing.T) {
		builder := NewMessageBuilder()
		params := map[string]interface{}{"type": "update"}

		message := builder.Notification("notification").WithParams(params).Build()

		if message.JSONRPC != "2.0" {
			t.Errorf("Expected JSONRPC 2.0, got %s", message.JSONRPC)
		}

		if message.ID != nil {
			t.Errorf("Expected nil ID for notification, got %v", message.ID)
		}

		if message.Method != "notification" {
			t.Errorf("Expected method 'notification', got %s", message.Method)
		}
	})

	t.Run("BuildError", func(t *testing.T) {
		builder := NewMessageBuilder()
		errorData := map[string]interface{}{"details": "test error"}

		message := builder.Response(1).WithError(-32603, "Internal error", errorData).Build()

		if message.Error == nil {
			t.Error("Expected non-nil error")
		}

		if message.Error.Code != -32603 {
			t.Errorf("Expected error code -32603, got %d", message.Error.Code)
		}

		if message.Error.Message != "Internal error" {
			t.Errorf("Expected error message 'Internal error', got %s", message.Error.Message)
		}
	})

	t.Run("BuildAndValidate", func(t *testing.T) {
		builder := NewMessageBuilder()

		// Valid message
		validMessage, err := builder.Request(1, "test").BuildAndValidate()
		if err != nil {
			t.Errorf("Expected no validation error for valid message, got %v", err)
		}

		if validMessage == nil {
			t.Error("Expected non-nil message")
		}

		// Invalid message
		builder = NewMessageBuilder()
		_, err = builder.Response(1).BuildAndValidate() // Missing result or error
		if err == nil {
			t.Error("Expected validation error for invalid message")
		}
	})
}
