import { createServer, type Server } from 'node:http';
import { randomUUID } from 'node:crypto';
import { Client } from '../../src/client/index.js';
import { StreamableHTTPClientTransport } from '../../src/client/streamableHttp.js';
import { McpServer } from '../../src/server/mcp.js';
import { StreamableHTTPServerTransport } from '../../src/server/streamableHttp.js';
import {
    CallToolResultSchema,
    ListToolsResultSchema,
    ListResourcesResultSchema,
    ListPromptsResultSchema,
    LATEST_PROTOCOL_VERSION
} from '../../src/types.js';
import { zodTestMatrix, type ZodMatrixEntry } from '../../src/__fixtures__/zodTestMatrix.js';
import { listenOnRandomPort } from '../helpers/http.js';

describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
    const { z } = entry;
    describe('Streamable HTTP Transport Session Management', () => {
        /**
         * Helper to create and configure a fresh McpServer instance with standard
         * resources, prompts, and tools for testing.
         */
        function createMcpServer(): McpServer {
            const mcpServer = new McpServer(
                { name: 'test-server', version: '1.0.0' },
                {
                    capabilities: {
                        logging: {},
                        tools: {},
                        resources: {},
                        prompts: {}
                    }
                }
            );

            // Add a simple resource
            mcpServer.resource('test-resource', '/test', { description: 'A test resource' }, async () => ({
                contents: [
                    {
                        uri: '/test',
                        text: 'This is a test resource content'
                    }
                ]
            }));

            mcpServer.prompt('test-prompt', 'A test prompt', async () => ({
                messages: [
                    {
                        role: 'user',
                        content: {
                            type: 'text',
                            text: 'This is a test prompt'
                        }
                    }
                ]
            }));

            mcpServer.tool(
                'greet',
                'A simple greeting tool',
                {
                    name: z.string().describe('Name to greet').default('World')
                },
                async ({ name }) => {
                    return {
                        content: [{ type: 'text', text: `Hello, ${name}!` }]
                    };
                }
            );

            return mcpServer;
        }

        // Function to set up the server with optional session management
        async function setupServer(withSessionManagement: boolean): Promise<{
            server: Server;
            mcpServer?: McpServer;
            serverTransport?: StreamableHTTPServerTransport;
            baseUrl: URL;
        }> {
            const server: Server = createServer();

            if (withSessionManagement) {
                // Stateful mode: single transport + server for the session
                const mcpServer = createMcpServer();
                const serverTransport = new StreamableHTTPServerTransport({
                    sessionIdGenerator: () => randomUUID()
                });

                await mcpServer.connect(serverTransport);

                server.on('request', async (req, res) => {
                    await serverTransport.handleRequest(req, res);
                });

                // Start the server on a random port
                const baseUrl = await listenOnRandomPort(server);

                return { server, mcpServer, serverTransport, baseUrl };
            } else {
                // Stateless mode: create a fresh transport + server per request
                // to comply with the guard that stateless transports cannot be reused.
                server.on('request', async (req, res) => {
                    const mcpServer = createMcpServer();
                    const serverTransport = new StreamableHTTPServerTransport({
                        sessionIdGenerator: undefined
                    });
                    await mcpServer.connect(serverTransport);
                    await serverTransport.handleRequest(req, res);
                    // Close the per-request mcpServer after handling to avoid leaks
                    await mcpServer.close();
                });

                // Start the server on a random port
                const baseUrl = await listenOnRandomPort(server);

                return { server, baseUrl };
            }
        }

        describe('Stateless Mode', () => {
            let server: Server;
            let baseUrl: URL;

            beforeEach(async () => {
                const setup = await setupServer(false);
                server = setup.server;
                baseUrl = setup.baseUrl;
            });

            afterEach(async () => {
                server.close();
            });

            it('should support multiple client connections', async () => {
                // Create and connect a client
                const client1 = new Client({
                    name: 'test-client',
                    version: '1.0.0'
                });

                const transport1 = new StreamableHTTPClientTransport(baseUrl);
                await client1.connect(transport1);

                // Verify that no session ID was set
                expect(transport1.sessionId).toBeUndefined();

                // List available tools
                await client1.request(
                    {
                        method: 'tools/list',
                        params: {}
                    },
                    ListToolsResultSchema
                );

                const client2 = new Client({
                    name: 'test-client',
                    version: '1.0.0'
                });

                const transport2 = new StreamableHTTPClientTransport(baseUrl);
                await client2.connect(transport2);

                // Verify that no session ID was set
                expect(transport2.sessionId).toBeUndefined();

                // List available tools
                await client2.request(
                    {
                        method: 'tools/list',
                        params: {}
                    },
                    ListToolsResultSchema
                );
            });
            it('should operate without session management', async () => {
                // Create and connect a client
                const client = new Client({
                    name: 'test-client',
                    version: '1.0.0'
                });

                const transport = new StreamableHTTPClientTransport(baseUrl);
                await client.connect(transport);

                // Verify that no session ID was set
                expect(transport.sessionId).toBeUndefined();

                // List available tools
                const toolsResult = await client.request(
                    {
                        method: 'tools/list',
                        params: {}
                    },
                    ListToolsResultSchema
                );

                // Verify tools are accessible
                expect(toolsResult.tools).toContainEqual(
                    expect.objectContaining({
                        name: 'greet'
                    })
                );

                // List available resources
                const resourcesResult = await client.request(
                    {
                        method: 'resources/list',
                        params: {}
                    },
                    ListResourcesResultSchema
                );

                // Verify resources result structure
                expect(resourcesResult).toHaveProperty('resources');

                // List available prompts
                const promptsResult = await client.request(
                    {
                        method: 'prompts/list',
                        params: {}
                    },
                    ListPromptsResultSchema
                );

                // Verify prompts result structure
                expect(promptsResult).toHaveProperty('prompts');
                expect(promptsResult.prompts).toContainEqual(
                    expect.objectContaining({
                        name: 'test-prompt'
                    })
                );

                // Call the greeting tool
                const greetingResult = await client.request(
                    {
                        method: 'tools/call',
                        params: {
                            name: 'greet',
                            arguments: {
                                name: 'Stateless Transport'
                            }
                        }
                    },
                    CallToolResultSchema
                );

                // Verify tool result
                expect(greetingResult.content).toEqual([{ type: 'text', text: 'Hello, Stateless Transport!' }]);

                // Clean up
                await transport.close();
            });

            it('should set protocol version after connecting', async () => {
                // Create and connect a client
                const client = new Client({
                    name: 'test-client',
                    version: '1.0.0'
                });

                const transport = new StreamableHTTPClientTransport(baseUrl);

                // Verify protocol version is not set before connecting
                expect(transport.protocolVersion).toBeUndefined();

                await client.connect(transport);

                // Verify protocol version is set after connecting
                expect(transport.protocolVersion).toBe(LATEST_PROTOCOL_VERSION);

                // Clean up
                await transport.close();
            });
        });

        describe('Stateful Mode', () => {
            let server: Server;
            let mcpServer: McpServer;
            let serverTransport: StreamableHTTPServerTransport;
            let baseUrl: URL;

            beforeEach(async () => {
                const setup = await setupServer(true);
                server = setup.server;
                mcpServer = setup.mcpServer!;
                serverTransport = setup.serverTransport!;
                baseUrl = setup.baseUrl;
            });

            afterEach(async () => {
                // Clean up resources
                await mcpServer.close().catch(() => {});
                await serverTransport.close().catch(() => {});
                server.close();
            });

            it('should operate with session management', async () => {
                // Create and connect a client
                const client = new Client({
                    name: 'test-client',
                    version: '1.0.0'
                });

                const transport = new StreamableHTTPClientTransport(baseUrl);
                await client.connect(transport);

                // Verify that a session ID was set
                expect(transport.sessionId).toBeDefined();
                expect(typeof transport.sessionId).toBe('string');

                // List available tools
                const toolsResult = await client.request(
                    {
                        method: 'tools/list',
                        params: {}
                    },
                    ListToolsResultSchema
                );

                // Verify tools are accessible
                expect(toolsResult.tools).toContainEqual(
                    expect.objectContaining({
                        name: 'greet'
                    })
                );

                // List available resources
                const resourcesResult = await client.request(
                    {
                        method: 'resources/list',
                        params: {}
                    },
                    ListResourcesResultSchema
                );

                // Verify resources result structure
                expect(resourcesResult).toHaveProperty('resources');

                // List available prompts
                const promptsResult = await client.request(
                    {
                        method: 'prompts/list',
                        params: {}
                    },
                    ListPromptsResultSchema
                );

                // Verify prompts result structure
                expect(promptsResult).toHaveProperty('prompts');
                expect(promptsResult.prompts).toContainEqual(
                    expect.objectContaining({
                        name: 'test-prompt'
                    })
                );

                // Call the greeting tool
                const greetingResult = await client.request(
                    {
                        method: 'tools/call',
                        params: {
                            name: 'greet',
                            arguments: {
                                name: 'Stateful Transport'
                            }
                        }
                    },
                    CallToolResultSchema
                );

                // Verify tool result
                expect(greetingResult.content).toEqual([{ type: 'text', text: 'Hello, Stateful Transport!' }]);

                // Clean up
                await transport.close();
            });
        });
    });
});