module mcp

import net.http
import time

fn test_server_routes_initialize_and_registered_features() {
    mut server := new_server(
        name:         'test-server'
        version:      '1.2.3'
        instructions: 'Be precise.'
    )
    server.add_tool(Tool{
        name:        'say_hello'
        description: 'Returns a greeting'
    }, fn (ctx Context, arguments string) !ToolResult {
        assert ctx.session_id == stdio_session_id
        assert ctx.transport == .stdio
        assert arguments == '{"name":"V"}'
        return tool_text_result('Hello, V!')
    })!
    server.add_resource(Resource{
        uri:       'resource://guide'
        name:      'guide'
        mime_type: 'text/plain'
    }, fn (_ Context, uri string) !ReadResourceResult {
        return ReadResourceResult{
            contents: [
                ResourceContents{
                    uri:       uri
                    mime_type: 'text/plain'
                    text:      'guide contents'
                },
            ]
        }
    })!
    server.add_resource_template(ResourceTemplate{
        uri_template: 'resource://docs/{slug}'
        name:         'docs'
    })!
    server.add_prompt(Prompt{
        name:        'review'
        description: 'Review some code'
        arguments:   [
            PromptArgument{
                name:     'code'
                required: true
            },
        ]
    }, fn (_ Context, arguments string) !GetPromptResult {
        assert arguments == '{"code":"fn main() {}"}'
        return GetPromptResult{
            description: 'Review prompt'
            messages:    [
                prompt_text_message('user', 'Review this code'),
            ]
        }
    })!

    init_request := Request{
        id:     encode_id(1)
        method: 'initialize'
        params: encode_initialize_params(InitializeParams{
            protocol_version: protocol_version
            capabilities:     '{"roots":{}}'
            client_info:      Implementation{
                name:    'test-client'
                version: '0.1.0'
            }
        })
    }
    init_dispatch := server.dispatch_message(init_request.encode(), stdio_session_id, .stdio)!
    assert init_dispatch.has_response
    init_response := decode_response(init_dispatch.response)!
    init_result := init_response.decode_result[InitializeResult]()!
    assert init_result.server_info.name == 'test-server'
    assert init_result.instructions == 'Be precise.'
    assert init_result.capabilities == '{"tools":{},"resources":{},"prompts":{}}'

    blocked_dispatch := server.dispatch_message(new_request(2, 'tools/list', empty).encode(),
        stdio_session_id, .stdio)!
    blocked_response := decode_response(blocked_dispatch.response)!
    assert blocked_response.error.code == server_not_initialized.code

    initialized := server.dispatch_message(new_notification('notifications/initialized', empty).encode(),
        stdio_session_id, .stdio)!
    assert !initialized.has_response

    tools_list := server.dispatch_message(new_request(3, 'tools/list', empty).encode(),
        stdio_session_id, .stdio)!
    tools_result := decode_response(tools_list.response)!.decode_result[ListToolsResult]()!
    assert tools_result.tools.len == 1
    assert tools_result.tools[0].name == 'say_hello'

    tool_call := server.dispatch_message(Request{
        id:     encode_id(4)
        method: 'tools/call'
        params: '{"name":"say_hello","arguments":{"name":"V"}}'
    }.encode(), stdio_session_id, .stdio)!
    tool_result := decode_response(tool_call.response)!.decode_result[ToolResult]()!
    assert tool_result.content.contains('Hello, V!')
    assert !tool_result.is_error

    resource_list := server.dispatch_message(new_request(5, 'resources/list', empty).encode(),
        stdio_session_id, .stdio)!
    resource_list_result :=
        decode_response(resource_list.response)!.decode_result[ListResourcesResult]()!
    assert resource_list_result.resources.len == 1
    assert resource_list_result.resources[0].uri == 'resource://guide'

    resource_templates := server.dispatch_message(new_request(6, 'resources/templates/list', empty).encode(),
        stdio_session_id, .stdio)!
    resource_template_result :=
        decode_response(resource_templates.response)!.decode_result[ListResourceTemplatesResult]()!
    assert resource_template_result.resource_templates.len == 1
    assert resource_template_result.resource_templates[0].uri_template == 'resource://docs/{slug}'

    resource_read := server.dispatch_message(new_request(7, 'resources/read', ReadResourceParams{
        uri: 'resource://guide'
    }).encode(), stdio_session_id, .stdio)!
    resource_read_result :=
        decode_response(resource_read.response)!.decode_result[ReadResourceResult]()!
    assert resource_read_result.contents.len == 1
    assert resource_read_result.contents[0].text == 'guide contents'

    prompts_list := server.dispatch_message(new_request(8, 'prompts/list', empty).encode(),
        stdio_session_id, .stdio)!
    prompts_list_result :=
        decode_response(prompts_list.response)!.decode_result[ListPromptsResult]()!
    assert prompts_list_result.prompts.len == 1
    assert prompts_list_result.prompts[0].name == 'review'

    prompt_get := server.dispatch_message(Request{
        id:     encode_id(9)
        method: 'prompts/get'
        params: '{"name":"review","arguments":{"code":"fn main() {}"}}'
    }.encode(), stdio_session_id, .stdio)!
    prompt_result := decode_response(prompt_get.response)!.decode_result[GetPromptResult]()!
    assert prompt_result.messages.len == 1
    assert prompt_result.messages[0].role == 'user'

    ping := server.dispatch_message(new_request(10, 'ping', empty).encode(), stdio_session_id,
        .stdio)!
    ping_result := decode_response(ping.response)!.decode_result[EmptyObject]()!
    assert ping_result == empty_object
}

fn test_server_http_sessions_and_delete() {
    mut server_value := new_server(
        name:    'http-server'
        version: '0.0.1'
    )
    mut server := &server_value
    server.add_tool(Tool{
        name: 'ping_tool'
    }, fn (_ Context, _ string) !ToolResult {
        return tool_text_result('pong')
    })!

    server_thread := spawn server.serve_http('127.0.0.1:0')
    server.wait_till_running(max_retries: 200, retry_period_ms: 10)!
    time.sleep(20 * time.millisecond)
    addr := server.http_server.addr
    url := 'http://${addr}/mcp'

    mut init_header := http.new_header(http.HeaderConfig{
        key:   .content_type
        value: 'application/json'
    }, http.HeaderConfig{
        key:   .accept
        value: 'application/json'
    })
    init_response := http.fetch(
        method: .post
        url:    url
        data:   Request{
            id:     encode_id(1)
            method: 'initialize'
            params: encode_initialize_params(InitializeParams{
                protocol_version: protocol_version
                capabilities:     '{}'
                client_info:      Implementation{
                    name:    'http-client'
                    version: '0.1.0'
                }
            })
        }.encode()
        header: init_header
    )!
    assert init_response.status_code == 200
    session_id := init_response.header.get_custom(mcp_session_id_header) or {
        assert false
        return
    }
    assert session_id != ''

    mut notification_header := init_header
    notification_header.set_custom(mcp_session_id_header, session_id)!
    notification_response := http.fetch(
        method: .post
        url:    url
        data:   new_notification('notifications/initialized', empty).encode()
        header: notification_header
    )!
    assert notification_response.status_code == 202

    mut list_header := notification_header
    list_header.set(.accept, 'text/event-stream')
    list_response := http.fetch(
        method: .post
        url:    url
        data:   new_request(2, 'tools/list', empty).encode()
        header: list_header
    )!
    assert list_response.status_code == 200
    assert list_response.header.get(.content_type)!.starts_with(event_stream_content_type)
    list_messages := parse_sse_messages(list_response.body)!
    assert list_messages.len == 1
    list_result := decode_response(list_messages[0])!.decode_result[ListToolsResult]()!
    assert list_result.tools.len == 1
    assert list_result.tools[0].name == 'ping_tool'

    mut delete_header := http.new_header()
    delete_header.set_custom(mcp_session_id_header, session_id)!
    delete_response := http.fetch(
        method: .delete
        url:    url
        header: delete_header
    )!
    assert delete_response.status_code == 200

    mut stale_header := init_header
    stale_header.set_custom(mcp_session_id_header, session_id)!
    stale_response := http.fetch(
        method: .post
        url:    url
        data:   new_request(3, 'tools/list', empty).encode()
        header: stale_header
    )!
    assert stale_response.status_code == 404
    server.close()
    server_thread.wait() or {}
}