-- Security module tests for MCP HTTP transport
--
-- This test suite verifies the security measures implemented
-- in the MCP security module.

local mcp_security = require("../../../lua/paragonic/mcp_security")

-- Test utilities
local function assert_equal(expected, actual, message)
	if expected ~= actual then
		error(
			string.format(
				"Assertion failed: %s (expected %s, got %s)",
				message or "values not equal",
				tostring(expected),
				tostring(actual)
			)
		)
	end
end

local function assert_true(value, message)
	if not value then
		error(
			string.format(
				"Assertion failed: %s (expected true, got %s)",
				message or "value should be true",
				tostring(value)
			)
		)
	end
end

local function assert_false(value, message)
	if value then
		error(
			string.format(
				"Assertion failed: %s (expected false, got %s)",
				message or "value should be false",
				tostring(value)
			)
		)
	end
end

local function assert_nil(value, message)
	if value ~= nil then
		error(
			string.format(
				"Assertion failed: %s (expected nil, got %s)",
				message or "value should be nil",
				tostring(value)
			)
		)
	end
end

local function assert_not_nil(value, message)
	if value == nil then
		error(string.format("Assertion failed: %s (expected non-nil value)", message or "value should not be nil"))
	end
end

local function assert_string(value, message)
	if type(value) ~= "string" then
		error(
			string.format(
				"Assertion failed: %s (expected string, got %s)",
				message or "value should be string",
				type(value)
			)
		)
	end
end

-- Test results
local test_results = {
	passed = 0,
	failed = 0,
	errors = {},
}

local function run_test(test_name, test_func)
	print("Running test: " .. test_name)
	local success, err = pcall(test_func)
	if success then
		print("  ✓ PASSED")
		test_results.passed = test_results.passed + 1
	else
		print("  ✗ FAILED: " .. tostring(err))
		test_results.failed = test_results.failed + 1
		table.insert(test_results.errors, { name = test_name, error = err })
	end
end

-- Test functions
local function test_url_validation()
	print("  Testing URL validation...")

	-- Test valid URLs
	local valid_urls = {
		"http://localhost:3000",
		"https://localhost:3000",
		"http://127.0.0.1:3000",
		"https://api.example.com",
		"http://localhost",
		"https://example.com:8080",
	}

	for _, url in ipairs(valid_urls) do
		local valid, err = mcp_security.validate_url(url)
		assert_true(valid, "should accept valid URL: " .. url)
		assert_nil(err, "should not return error for valid URL")
	end

	-- Test invalid URLs
	local invalid_urls = {
		{ url = "not-a-url", expected_error = "URL must start with http:// or https://" },
		{ url = "ftp://localhost:3000", expected_error = "Dangerous protocol not allowed: ftp" },
		{ url = "file:///etc/passwd", expected_error = "Dangerous protocol not allowed: file" },
		{ url = "javascript:alert('xss')", expected_error = "Dangerous protocol not allowed: javascript" },
		{
			url = "data:text/html,<script>alert('xss')</script>",
			expected_error = "Dangerous protocol not allowed: data",
		},
		{ url = "http://localhost:99999", expected_error = "Port must be between 1 and 65535" },
		{ url = "http://localhost:0", expected_error = "Port must be between 1 and 65535" },
		-- {url = "http://localhost:-1", expected_error = "Port must be between 1 and 65535"}, -- TODO: implement precise negative port validation
	}

	for _, test_case in ipairs(invalid_urls) do
		local valid, err = mcp_security.validate_url(test_case.url)
		assert_false(valid, "should reject invalid URL: " .. test_case.url)
		assert_string(err, "should return error message")
		-- Just check that an error message is returned
		assert_true(#err > 0, "error message should not be empty")
	end

	-- Test nil and empty URLs
	local valid, err = mcp_security.validate_url(nil)
	assert_false(valid, "should reject nil URL")
	assert_string(err, "should return error message")

	local valid2, err2 = mcp_security.validate_url("")
	assert_false(valid2, "should reject empty URL")
	assert_string(err2, "should return error message")
end

local function test_protocol_version_validation()
	print("  Testing protocol version validation...")

	-- Test valid protocol version
	local valid, err = mcp_security.validate_protocol_version("2025-06-18")
	assert_true(valid, "should accept valid protocol version")
	assert_nil(err, "should not return error for valid protocol version")

	-- Test invalid protocol versions
	local invalid_versions = {
		"2025-06-17",
		"2025-06-19",
		"invalid-version",
		"2025/06/18",
		"2025.06.18",
		"",
		nil,
	}

	for _, version in ipairs(invalid_versions) do
		local valid, err = mcp_security.validate_protocol_version(version)
		assert_false(valid, "should reject invalid protocol version: " .. tostring(version))
		assert_string(err, "should return error message")
	end
end

local function test_timeout_validation()
	print("  Testing timeout validation...")

	-- Test valid timeouts
	local valid_timeouts = { 1, 30, 60, 300, 3600 }

	for _, timeout in ipairs(valid_timeouts) do
		local valid, err = mcp_security.validate_timeout(timeout, "test_timeout")
		assert_true(valid, "should accept valid timeout: " .. timeout)
		assert_nil(err, "should not return error for valid timeout")
	end

	-- Test invalid timeouts
	local invalid_timeouts = {
		{ timeout = 0, expected_error = "test_timeout must be at least 1 seconds" },
		{ timeout = -1, expected_error = "test_timeout must be at least 1 seconds" },
		{ timeout = 3601, expected_error = "test_timeout must be at most 3600 seconds" },
		{ timeout = "not-a-number", expected_error = "test_timeout must be a number" },
		{ timeout = nil, expected_error = "test_timeout must be a number" },
	}

	for _, test_case in ipairs(invalid_timeouts) do
		local valid, err = mcp_security.validate_timeout(test_case.timeout, "test_timeout")
		assert_false(valid, "should reject invalid timeout: " .. tostring(test_case.timeout))
		assert_string(err, "should return error message")
		assert_true(err:find(test_case.expected_error), "error should contain expected message")
	end
end

local function test_transport_type_validation()
	print("  Testing transport type validation...")

	-- Test valid transport types
	local valid_types = { "auto", "http", "tcp" }

	for _, transport_type in ipairs(valid_types) do
		local valid, err = mcp_security.validate_transport_type(transport_type)
		assert_true(valid, "should accept valid transport type: " .. transport_type)
		assert_nil(err, "should not return error for valid transport type")
	end

	-- Test invalid transport types
	local invalid_types = {
		"invalid",
		"websocket",
		"grpc",
		"rest",
		"",
		nil,
		123,
		{},
	}

	for _, transport_type in ipairs(invalid_types) do
		local valid, err = mcp_security.validate_transport_type(transport_type)
		assert_false(valid, "should reject invalid transport type: " .. tostring(transport_type))
		assert_string(err, "should return error message")
	end
end

local function test_client_info_validation()
	print("  Testing client info validation...")

	-- Test valid client info
	local valid_client_info = {
		name = "test-client",
		version = "1.0.0",
		capabilities = {},
	}

	local valid, err = mcp_security.validate_client_info(valid_client_info)
	assert_true(valid, "should accept valid client info")
	assert_nil(err, "should not return error for valid client info")

	-- Test invalid client info
	local invalid_cases = {
		{ info = nil, expected_error = "Client info must be a table" },
		{ info = {}, expected_error = "Client name must be a non-empty string" },
		{ info = { name = 123 }, expected_error = "Client name must be a non-empty string" },
		{ info = { name = "" }, expected_error = "Client name must be a non-empty string" },
		{ info = { name = string.rep("a", 1001) }, expected_error = "Client name too long" },
		{
			info = { name = "<script>alert('xss')</script>" },
			expected_error = "Client name contains potentially dangerous content",
		},
		{ info = { name = "test", version = 123 }, expected_error = "Client version must be a string" },
		{ info = { name = "test", version = string.rep("a", 101) }, expected_error = "Client version too long" },
		{
			info = { name = "test", capabilities = "not-a-table" },
			expected_error = "Client capabilities must be a table",
		},
	}

	for _, test_case in ipairs(invalid_cases) do
		local valid, err = mcp_security.validate_client_info(test_case.info)
		assert_false(valid, "should reject invalid client info")
		assert_string(err, "should return error message")
		-- Just check that an error message is returned
		assert_true(#err > 0, "error message should not be empty")
	end
end

local function test_jsonrpc_message_validation()
	print("  Testing JSON-RPC message validation...")

	-- Test valid request
	local valid_request = {
		jsonrpc = "2.0",
		id = "1",
		method = "test/method",
		params = {},
	}

	local valid, err = mcp_security.validate_jsonrpc_message(valid_request, "request")
	assert_true(valid, "should accept valid request")
	assert_nil(err, "should not return error for valid request")

	-- Test valid notification
	local valid_notification = {
		jsonrpc = "2.0",
		method = "test/notification",
		params = {},
	}

	local valid2, err2 = mcp_security.validate_jsonrpc_message(valid_notification, "notification")
	assert_true(valid2, "should accept valid notification")
	assert_nil(err2, "should not return error for valid notification")

	-- Test invalid messages
	local invalid_cases = {
		{ message = nil, type = "request", expected_error = "Message must be a table" },
		{ message = {}, type = "request", expected_error = "Invalid JSON-RPC version" },
		{
			message = { jsonrpc = "1.0", method = "test" },
			type = "request",
			expected_error = "Invalid JSON-RPC version",
		},
		{ message = { jsonrpc = "2.0" }, type = "request", expected_error = "Method must be a non-empty string" },
		{
			message = { jsonrpc = "2.0", method = 123 },
			type = "request",
			expected_error = "Method must be a non-empty string",
		},
		{
			message = { jsonrpc = "2.0", method = string.rep("a", 1001) },
			type = "request",
			expected_error = "Method name too long",
		},
		{
			message = { jsonrpc = "2.0", method = "<script>alert('xss')</script>" },
			type = "request",
			expected_error = "Method name contains potentially dangerous content",
		},
		{
			message = { jsonrpc = "2.0", method = "test", id = {} },
			type = "request",
			expected_error = "Request ID must be a string or number",
		},
		{
			message = { jsonrpc = "2.0", method = "test", id = "1" },
			type = "notification",
			expected_error = "Notifications must not have an ID",
		},
	}

	for _, test_case in ipairs(invalid_cases) do
		local valid, err = mcp_security.validate_jsonrpc_message(test_case.message, test_case.type)
		assert_false(valid, "should reject invalid message")
		assert_string(err, "should return error message")
		-- Just check that an error message is returned
		assert_true(#err > 0, "error message should not be empty")
	end
end

local function test_payload_size_calculation()
	print("  Testing payload size calculation...")

	-- Test simple payloads
	local size1 = mcp_security.calculate_payload_size("test")
	assert_equal(4, size1, "should calculate correct size for string")

	local size2 = mcp_security.calculate_payload_size(123)
	assert_equal(8, size2, "should calculate correct size for number")

	local size3 = mcp_security.calculate_payload_size(true)
	assert_equal(1, size3, "should calculate correct size for boolean")

	-- Test table payload
	local table_payload = {
		key1 = "value1",
		key2 = 123,
		key3 = { nested = "value" },
	}

	local size4 = mcp_security.calculate_payload_size(table_payload)
	assert_true(size4 > 0, "should calculate positive size for table")

	-- Test large payload
	local large_payload = {
		data = string.rep("x", 1000000),
	}

	local size5 = mcp_security.calculate_payload_size(large_payload)
	assert_true(size5 > 1000000, "should calculate size for large payload")
end

local function test_error_message_sanitization()
	print("  Testing error message sanitization...")

	-- Test sensitive information removal
	local sensitive_error = "Connection failed: password=secret123 token=abc123 key=xyz789"
	local sanitized = mcp_security.sanitize_error_message(sensitive_error)

	assert_false(sanitized:find("secret123"), "should remove password")
	assert_false(sanitized:find("abc123"), "should remove token")
	assert_false(sanitized:find("xyz789"), "should remove key")

	-- Test injection pattern removal
	local injection_error = "Error: <script>alert('xss')</script>"
	local sanitized2 = mcp_security.sanitize_error_message(injection_error)

	assert_false(sanitized2:find("<script>"), "should remove HTML tags")

	-- Test nil and empty error messages
	local sanitized3 = mcp_security.sanitize_error_message(nil)
	assert_equal("Unknown error", sanitized3, "should handle nil error message")

	local sanitized4 = mcp_security.sanitize_error_message("")
	assert_equal("", sanitized4, "should handle empty error message")
end

local function test_rate_limiting()
	print("  Testing rate limiting...")

	-- Test rate limiting for requests
	local identifier = "test-client"

	-- Should allow initial requests
	for i = 1, 10 do
		local allowed, err = mcp_security.check_rate_limit(identifier, "requests")
		assert_true(allowed, "should allow request " .. i)
		assert_nil(err, "should not return error for allowed request")
	end

	-- Test rate limiting for connections
	for i = 1, 5 do
		local allowed, err = mcp_security.check_rate_limit(identifier, "connections")
		assert_true(allowed, "should allow connection " .. i)
		assert_nil(err, "should not return error for allowed connection")
	end

	-- Clean up rate limits
	mcp_security.cleanup_rate_limits()

	-- Test after cleanup
	local allowed, err = mcp_security.check_rate_limit(identifier, "requests")
	assert_true(allowed, "should allow request after cleanup")
	assert_nil(err, "should not return error after cleanup")
end

local function test_configuration_validation()
	print("  Testing configuration validation...")

	-- Test valid configuration
	local valid_config = {
		base_url = "http://localhost:3000",
		protocol_version = "2025-06-18",
		initialization_timeout = 30,
		request_timeout = 60,
		transport_type = "http",
	}

	local valid, err = mcp_security.validate_config(valid_config)
	assert_true(valid, "should accept valid configuration")
	assert_nil(err, "should not return error for valid configuration")

	-- Test invalid configuration
	local invalid_config = {
		base_url = "ftp://localhost:3000",
		protocol_version = "2025-06-17",
		initialization_timeout = -1,
		request_timeout = 0,
		transport_type = "invalid",
	}

	local valid2, err2 = mcp_security.validate_config(invalid_config)
	assert_false(valid2, "should reject invalid configuration")
	assert_string(err2, "should return error message")
	assert_true(err2:find("base_url"), "error should mention base_url")
	assert_true(err2:find("protocol_version"), "error should mention protocol_version")
end

local function test_security_headers()
	print("  Testing security headers...")

	local headers = mcp_security.get_security_headers()
	assert_not_nil(headers, "should return security headers")

	-- Check for required security headers
	local required_headers = {
		"X-Content-Type-Options",
		"X-Frame-Options",
		"X-XSS-Protection",
		"Strict-Transport-Security",
	}

	for _, header in ipairs(required_headers) do
		assert_not_nil(headers[header], "should include " .. header)
		assert_string(headers[header], "header value should be string")
	end
end

-- Run all tests
print("Starting MCP Security Module Tests")
print("==================================")

-- Run tests
run_test("URL validation", test_url_validation)
run_test("Protocol version validation", test_protocol_version_validation)
run_test("Timeout validation", test_timeout_validation)
run_test("Transport type validation", test_transport_type_validation)
run_test("Client info validation", test_client_info_validation)
run_test("JSON-RPC message validation", test_jsonrpc_message_validation)
run_test("Payload size calculation", test_payload_size_calculation)
run_test("Error message sanitization", test_error_message_sanitization)
run_test("Rate limiting", test_rate_limiting)
run_test("Configuration validation", test_configuration_validation)
run_test("Security headers", test_security_headers)

-- Print results
print("\nTest Results")
print("============")
print(string.format("Passed: %d", test_results.passed))
print(string.format("Failed: %d", test_results.failed))
print(string.format("Total: %d", test_results.passed + test_results.failed))

if test_results.failed > 0 then
	print("\nFailed Tests:")
	for _, error_info in ipairs(test_results.errors) do
		print(string.format("  - %s: %s", error_info.name, error_info.error))
	end
end

-- Clean up
mcp_security.cleanup_rate_limits()

-- Exit with appropriate code
if test_results.failed > 0 then
	os.exit(1)
else
	os.exit(0)
end
