module mcp

import json
import net.http
import os
import time

pub const jsonrpc_version = '2.0'
pub const protocol_version = '2025-11-25'
pub const parse_error = ResponseError{
    code:    -32700
    message: 'Invalid JSON.'
}
pub const invalid_request = ResponseError{
    code:    -32600
    message: 'Invalid request.'
}
pub const method_not_found = ResponseError{
    code:    -32601
    message: 'Method not found.'
}
pub const invalid_params = ResponseError{
    code:    -32602
    message: 'Invalid params.'
}
pub const internal_error = ResponseError{
    code:    -32603
    message: 'Internal error.'
}
pub const server_not_initialized = ResponseError{
    code:    -32002
    message: 'Server not initialized.'
}

const default_content_type = 'application/json'
const streamable_http_accept = 'application/json, text/event-stream'
const content_length_header = 'Content-Length'
const mcp_session_id_header = 'Mcp-Session-Id'
const process_poll_interval = 5 * time.millisecond
const default_client_name = 'v.mcp'
const default_client_version = 'dev'

// ResponseError is the JSON-RPC error payload used by MCP responses.
pub struct ResponseError {
pub:
    code    int
    message string
    data    string @[raw]
}

// code returns the JSON-RPC error code.
pub fn (err ResponseError) code() int {
    return err.code
}

// msg returns the JSON-RPC error message.
pub fn (err ResponseError) msg() string {
    return err.message
}

// err casts the response error to `IError`.
pub fn (err ResponseError) err() IError {
    return IError(err)
}

// Null represents the JSON `null` literal.
pub struct Null {}

// str returns the JSON `null` literal.
pub fn (n Null) str() string {
    return 'null'
}

pub const null = Null{}

// Empty omits a JSON-RPC field when used with MCP helpers.
pub struct Empty {}

// str returns the empty string.
pub fn (e Empty) str() string {
    return ''
}

pub const empty = Empty{}

// EmptyObject encodes to an empty JSON object.
pub struct EmptyObject {}

// str returns the JSON empty object literal.
pub fn (e EmptyObject) str() string {
    return '{}'
}

pub const empty_object = EmptyObject{}

// Implementation identifies an MCP client or server implementation.
pub struct Implementation {
pub:
    name    string
    version string
}

// InitializeParams is the typed payload for the `initialize` request.
pub struct InitializeParams {
pub:
    protocol_version string         @[json: protocolVersion]
    capabilities     string         @[raw]
    client_info      Implementation @[json: clientInfo]
}

// InitializeResult is the typed result returned by an MCP server after initialization.
pub struct InitializeResult {
pub:
    protocol_version string         @[json: protocolVersion]
    capabilities     string         @[raw]
    server_info      Implementation @[json: serverInfo]
    instructions     string
}

// Request is a JSON-RPC request message encoded for MCP.
pub struct Request {
pub:
    jsonrpc string = jsonrpc_version
    id      string @[raw]
    method  string
    params  string @[omitempty; raw]
}

// new_request constructs an MCP request with a typed id and params payload.
pub fn new_request[I, P](id I, method string, params P) Request {
    return Request{
        id:     encode_id(id)
        method: method
        params: encode_value(params)
    }
}

// encode serializes the request to JSON.
pub fn (req Request) encode() string {
    params_payload := if req.params.len == 0 { '' } else { ',"params":${req.params}' }
    id_payload := if req.id.len == 0 { null.str() } else { req.id }
    return '{"jsonrpc":"${jsonrpc_version}","id":${id_payload},"method":${json.encode(req.method)}${params_payload}}'
}

// decode_params decodes the raw request params into `T`.
pub fn (req Request) decode_params[T]() !T {
    return decode_value[T](req.params)
}

// Notification is a JSON-RPC notification encoded for MCP.
pub struct Notification {
pub:
    jsonrpc string = jsonrpc_version
    method  string
    params  string @[omitempty; raw]
}

// new_notification constructs an MCP notification with a typed params payload.
pub fn new_notification[P](method string, params P) Notification {
    return Notification{
        method: method
        params: encode_value(params)
    }
}

// encode serializes the notification to JSON.
pub fn (notification Notification) encode() string {
    params_payload := if notification.params.len == 0 {
        ''
    } else {
        ',"params":${notification.params}'
    }
    return '{"jsonrpc":"${jsonrpc_version}","method":${json.encode(notification.method)}${params_payload}}'
}

// decode_params decodes the raw notification params into `T`.
pub fn (notification Notification) decode_params[T]() !T {
    return decode_value[T](notification.params)
}

// Response is a JSON-RPC response message encoded for MCP.
pub struct Response {
pub:
    jsonrpc string = jsonrpc_version
    id      string @[raw]
    result  string @[raw]
    error   ResponseError
}

// new_response constructs an MCP response with a typed id and result payload.
pub fn new_response[I, R](id I, result R, err ResponseError) Response {
    return Response{
        id:     encode_id(id)
        result: if err.code != 0 { '' } else { encode_value(result) }
        error:  err
    }
}

// encode serializes the response to JSON.
pub fn (resp Response) encode() string {
    mut payload := '{"jsonrpc":"${jsonrpc_version}"'
    if resp.error.code != 0 {
        payload += ',"error":' + json.encode(resp.error)
    } else {
        result_payload := if resp.result.len == 0 { null.str() } else { resp.result }
        payload += ',"result":' + result_payload
    }
    id_payload := if resp.id.len == 0 { null.str() } else { resp.id }
    return payload + ',"id":${id_payload}}'
}

// decode_result decodes the response result into `T`.
pub fn (resp Response) decode_result[T]() !T {
    if resp.error.code != 0 {
        return resp.error.err()
    }
    return decode_value[T](resp.result)
}

// decode_request decodes a JSON payload into an MCP request.
pub fn decode_request(raw string) !Request {
    return json.decode(Request, raw) or { return err }
}

// decode_notification decodes a JSON payload into an MCP notification.
pub fn decode_notification(raw string) !Notification {
    return json.decode(Notification, raw) or { return err }
}

// decode_response decodes a JSON payload into an MCP response.
pub fn decode_response(raw string) !Response {
    return json.decode(Response, raw) or { return err }
}

struct MessageEnvelope {
    jsonrpc string
    id      string @[raw]
    method  string
    params  string @[raw]
    result  string @[raw]
    error   ResponseError
}

fn (env MessageEnvelope) encode() string {
    if env.method.len != 0 {
        if env.id.len == 0 || env.id == null.str() {
            return Notification{
                method: env.method
                params: env.params
            }.encode()
        }
        return Request{
            id:     env.id
            method: env.method
            params: env.params
        }.encode()
    }
    return Response{
        id:     env.id
        result: env.result
        error:  env.error
    }.encode()
}

fn decode_envelope(raw string) !MessageEnvelope {
    return json.decode(MessageEnvelope, raw) or { return err }
}

// Transport is the boundary between MCP messages and the wire format.
pub interface Transport {
mut:
    send(message string) !
    receive() !string
    close()
}

@[params]
pub struct ClientConfig {
pub mut:
    protocol_version string         = protocol_version
    client_info      Implementation = Implementation{
        name:    default_client_name
        version: default_client_version
    }
    capabilities     string = '{}'
    headers          map[string]string
}

pub struct Client {
mut:
    transport         Transport
    config            ClientConfig
    next_id           int = 1
    initialized       bool
    init_result       InitializeResult
    pending_responses map[string]Response
    notifications     []Notification
    server_requests   []Request
}

// new_client constructs an MCP client on top of a custom transport.
pub fn new_client(transport Transport, config ClientConfig) Client {
    return Client{
        transport:         transport
        config:            config
        pending_responses: map[string]Response{}
    }
}

// connect creates an MCP client for a streamable HTTP endpoint.
pub fn connect(url string) !Client {
    return connect_http(url, ClientConfig{})
}

// connect_http creates an MCP client for a streamable HTTP endpoint.
pub fn connect_http(url string, config ClientConfig) !Client {
    transport := new_http_transport(url, config)!
    return new_client(transport, config)
}

// connect_stdio creates an MCP client that talks to a local stdio server process.
pub fn connect_stdio(command string, args []string, config ClientConfig) !Client {
    transport := new_process_transport(command, args)!
    return new_client(transport, config)
}

// initialize starts the MCP initialization handshake using the client's config.
pub fn (mut c Client) initialize() !InitializeResult {
    return c.initialize_with_raw(c.config.capabilities, c.config.client_info)
}

// initialize_with starts the MCP initialization handshake using typed capabilities.
pub fn (mut c Client) initialize_with[X](capabilities X, client_info Implementation) !InitializeResult {
    return c.initialize_with_raw(encode_value(capabilities), client_info)
}

// send_request sends a typed request and waits for its response.
pub fn (mut c Client) send_request(request Request) !Response {
    if request.method == 'initialize' {
        return error('mcp.Client.initialize must be used for the MCP handshake')
    }
    c.ensure_initialized()!
    c.transport.send(request.encode())!
    return c.wait_for_response(request.id)
}

// request_message sends a method call and returns the raw MCP response.
pub fn (mut c Client) request_message[P](method string, params P) !Response {
    request := new_request(c.next_request_id(), method, params)
    return c.send_request(request)
}

// request sends a method call and decodes its result into `Result`.
pub fn (mut c Client) request[P, R](method string, params P) !R {
    response := c.request_message(method, params)!
    result := response.decode_result[R]()!
    return result
}

// send_notification sends a typed notification message.
pub fn (mut c Client) send_notification(notification Notification) ! {
    if notification.method == 'notifications/initialized' {
        return error('notifications/initialized is sent automatically after initialization')
    }
    c.ensure_initialized()!
    c.transport.send(notification.encode())!
}

// notify sends a method notification with a typed params payload.
pub fn (mut c Client) notify[P](method string, params P) ! {
    c.send_notification(new_notification(method, params))!
}

// take_notifications drains notifications queued while waiting for responses.
pub fn (mut c Client) take_notifications() []Notification {
    if c.notifications.len == 0 {
        return []Notification{}
    }
    drained := c.notifications.clone()
    c.notifications = []Notification{}
    return drained
}

// take_requests drains server initiated requests queued while waiting for responses.
pub fn (mut c Client) take_requests() []Request {
    if c.server_requests.len == 0 {
        return []Request{}
    }
    drained := c.server_requests.clone()
    c.server_requests = []Request{}
    return drained
}

// close releases the underlying transport.
pub fn (mut c Client) close() {
    c.transport.close()
}

fn (mut c Client) initialize_with_raw(capabilities string, client_info Implementation) !InitializeResult {
    if c.initialized {
        return c.init_result
    }
    c.config.capabilities = normalize_capabilities(capabilities)
    c.config.client_info = normalize_client_info(client_info)
    params := InitializeParams{
        protocol_version: normalize_protocol_version(c.config.protocol_version)
        capabilities:     c.config.capabilities
        client_info:      c.config.client_info
    }
    request := Request{
        id:     encode_id(c.next_request_id())
        method: 'initialize'
        params: encode_initialize_params(params)
    }
    c.transport.send(request.encode())!
    response := c.wait_for_response(request.id)!
    result := response.decode_result[InitializeResult]()!
    c.transport.send(new_notification('notifications/initialized', empty).encode())!
    c.initialized = true
    c.init_result = result
    return result
}

fn (mut c Client) ensure_initialized() ! {
    if !c.initialized {
        c.initialize()!
    }
}

fn (mut c Client) next_request_id() int {
    request_id := c.next_id
    c.next_id++
    return request_id
}

fn (mut c Client) wait_for_response(expected_id string) !Response {
    if expected_id in c.pending_responses {
        response := c.pending_responses[expected_id]
        c.pending_responses.delete(expected_id)
        return response
    }
    for {
        raw_message := c.transport.receive()!
        envelope := decode_envelope(raw_message)!
        if envelope.method.len != 0 {
            if envelope.id.len == 0 || envelope.id == null.str() {
                c.notifications << Notification{
                    method: envelope.method
                    params: envelope.params
                }
            } else {
                c.server_requests << Request{
                    id:     envelope.id
                    method: envelope.method
                    params: envelope.params
                }
            }
            continue
        }
        response := Response{
            id:     envelope.id
            result: envelope.result
            error:  envelope.error
        }
        if response.id == expected_id {
            return response
        }
        c.pending_responses[response.id] = response
    }
    return error('mcp: response loop exited unexpectedly')
}

fn encode_initialize_params(params InitializeParams) string {
    return '{"protocolVersion":${json.encode(params.protocol_version)},"capabilities":${normalize_capabilities(params.capabilities)},"clientInfo":${json.encode(params.client_info)}}'
}

struct NoFrameError {}

fn (err NoFrameError) msg() string {
    return 'no complete frame available'
}

fn (err NoFrameError) code() int {
    return 0
}

struct FrameExtraction {
    message   string
    remaining string
}

struct HttpTransport {
mut:
    url        string
    header     http.Header
    session_id string
    pending    []string
}

fn new_http_transport(url string, config ClientConfig) !HttpTransport {
    if url == '' {
        return error('mcp.connect_http: empty url')
    }
    if !url.starts_with('http://') && !url.starts_with('https://') {
        return error('mcp.connect_http: expected an http:// or https:// MCP endpoint')
    }
    mut header := http.new_header()
    header.set(.user_agent, default_client_name)
    if config.headers.len != 0 {
        header.add_custom_map(config.headers)!
    }
    return HttpTransport{
        url:    url
        header: header
    }
}

fn (mut transport HttpTransport) send(message string) ! {
    mut header := transport.header
    header.set(.content_type, default_content_type)
    header.set(.accept, streamable_http_accept)
    if transport.session_id != '' {
        header.set_custom(mcp_session_id_header, transport.session_id)!
    }
    response := http.fetch(
        method: .post
        url:    transport.url
        data:   message
        header: header
    )!
    if session_id := response.header.get_custom(mcp_session_id_header) {
        transport.session_id = session_id
    }
    messages := parse_http_response_messages(response)!
    if messages.len != 0 {
        transport.pending << messages
        return
    }
    if response.status_code >= 400 {
        return error('mcp.http: server returned HTTP ${response.status_code} without an MCP payload')
    }
}

fn (mut transport HttpTransport) receive() !string {
    if transport.pending.len == 0 {
        return error('mcp.http: no pending messages are available')
    }
    message := transport.pending[0]
    transport.pending = if transport.pending.len == 1 {
        []string{}
    } else {
        transport.pending[1..].clone()
    }
    return message
}

fn (mut transport HttpTransport) close() {
    if transport.session_id == '' {
        return
    }
    mut header := transport.header
    header.set_custom(mcp_session_id_header, transport.session_id) or { return }
    http.fetch(
        method: .delete
        url:    transport.url
        header: header
    ) or {}
    transport.session_id = ''
}

struct ProcessTransport {
mut:
    process &os.Process
    buffer  string
}

fn new_process_transport(command string, args []string) !ProcessTransport {
    if command == '' {
        return error('mcp.connect_stdio: empty command')
    }
    mut process := os.new_process(command)
    process.set_args(args)
    process.set_redirect_stdio()
    process.run()
    return ProcessTransport{
        process: process
    }
}

fn (mut transport ProcessTransport) send(message string) ! {
    transport.process.stdin_write(encode_framed_message(message))
}

fn (mut transport ProcessTransport) receive() !string {
    for {
        frame := try_extract_framed_message(transport.buffer) or {
            if err.msg() != NoFrameError{}.msg() {
                return err
            }
            FrameExtraction{}
        }
        if frame.message.len != 0 {
            transport.buffer = frame.remaining
            return frame.message
        }
        if transport.process.is_pending(.stdout) {
            chunk := transport.process.stdout_read()
            if chunk.len != 0 {
                transport.buffer += chunk
                continue
            }
        }
        if !transport.process.is_alive() {
            transport.buffer += transport.process.stdout_slurp()
            frame_after_exit := try_extract_framed_message(transport.buffer) or {
                if err.msg() != NoFrameError{}.msg() {
                    return err
                }
                FrameExtraction{}
            }
            if frame_after_exit.message.len != 0 {
                transport.buffer = frame_after_exit.remaining
                return frame_after_exit.message
            }
            stderr_output := transport.process.stderr_slurp().trim_space()
            if stderr_output.len != 0 {
                return error('mcp.stdio: process exited before a full MCP message was received: ${stderr_output}')
            }
            return error('mcp.stdio: process exited before a full MCP message was received')
        }
        time.sleep(process_poll_interval)
    }
    return error('mcp.stdio: receive loop exited unexpectedly')
}

fn (mut transport ProcessTransport) close() {
    if transport.process.is_alive() {
        transport.process.signal_term()
        for _ in 0 .. 20 {
            if !transport.process.is_alive() {
                break
            }
            time.sleep(10 * time.millisecond)
        }
        if transport.process.is_alive() {
            transport.process.signal_kill()
        } else if transport.process.status in [.running, .stopped] {
            transport.process.wait()
        }
    }
    transport.process.close()
}

fn parse_http_response_messages(response http.Response) ![]string {
    content_type := response.header.get(.content_type) or { '' }
    body := response.body.trim_space()
    if body.len == 0 {
        return []string{}
    }
    content_type_lower := content_type.to_lower()
    if content_type_lower.starts_with('application/json')
        || (content_type == '' && is_json_payload(body)) {
        return split_json_payloads(body)
    }
    if content_type_lower.starts_with('text/event-stream') {
        return parse_sse_messages(body)
    }
    return error('mcp.http: unsupported content type `${content_type}`')
}

fn split_json_payloads(body string) ![]string {
    trimmed := body.trim_space()
    if trimmed.len == 0 {
        return []string{}
    }
    if trimmed[0] != `[` {
        return [trimmed]
    }
    envelopes := json.decode([]MessageEnvelope, trimmed) or { return err }
    mut messages := []string{cap: envelopes.len}
    for envelope in envelopes {
        messages << envelope.encode()
    }
    return messages
}

fn parse_sse_messages(body string) ![]string {
    normalized := body.replace('\r\n', '\n').replace('\r', '\n')
    mut data_lines := []string{}
    mut messages := []string{}
    for line in normalized.split('\n') {
        if line.len == 0 {
            if data_lines.len != 0 {
                append_sse_payload(mut messages, data_lines.join('\n'))!
                data_lines = []string{}
            }
            continue
        }
        if line.starts_with(':') {
            continue
        }
        if line.starts_with('data:') {
            mut payload := line[5..]
            if payload.len != 0 && payload[0] == ` ` {
                payload = payload[1..]
            }
            data_lines << payload
        }
    }
    if data_lines.len != 0 {
        append_sse_payload(mut messages, data_lines.join('\n'))!
    }
    return messages
}

fn append_sse_payload(mut messages []string, payload string) ! {
    trimmed := payload.trim_space()
    if !is_json_payload(trimmed) {
        return
    }
    payloads := split_json_payloads(trimmed)!
    for item in payloads {
        messages << item
    }
}

fn is_json_payload(payload string) bool {
    trimmed := payload.trim_space()
    if trimmed.len == 0 {
        return false
    }
    return trimmed[0] == `{` || trimmed[0] == `[`
}

fn encode_framed_message(message string) string {
    return '${content_length_header}: ${message.len}\r\n\r\n${message}'
}

fn try_extract_framed_message(buffer string) !FrameExtraction {
    header_end := buffer.index('\r\n\r\n') or { return NoFrameError{} }
    header_text := buffer[..header_end]
    mut content_length := -1
    for line in header_text.split('\r\n') {
        if line.len == 0 {
            continue
        }
        parts := line.split_nth(':', 2)
        if parts.len != 2 {
            continue
        }
        if parts[0].trim_space().to_lower() != content_length_header.to_lower() {
            continue
        }
        length_text := parts[1].trim_space()
        parsed_length := length_text.int()
        if parsed_length == 0 && length_text != '0' {
            return error('mcp.stdio: invalid Content-Length header `${length_text}`')
        }
        content_length = parsed_length
        break
    }
    if content_length < 0 {
        return error('mcp.stdio: missing Content-Length header')
    }
    body_start := header_end + 4
    body_end := body_start + content_length
    if buffer.len < body_end {
        return NoFrameError{}
    }
    message := buffer[body_start..body_end]
    remaining := if body_end >= buffer.len { '' } else { buffer[body_end..] }
    return FrameExtraction{
        message:   message
        remaining: remaining
    }
}

fn encode_id[I](id I) string {
    return $if I is string {
        json.encode(id)
    } $else $if I is int {
        id.str()
    } $else {
        json.encode(id)
    }
}

fn encode_value[T](value T) string {
    return $if T is Empty {
        value.str()
    } $else $if T is EmptyObject {
        value.str()
    } $else $if T is Null {
        value.str()
    } $else $if T is string {
        json.encode(value)
    } $else {
        json.encode(value)
    }
}

fn decode_value[T](value string) !T {
    $if T is Empty {
        if value.len == 0 || value == null.str() {
            return Empty{}
        }
        return error('mcp: expected an empty payload, got `${value}`')
    } $else $if T is EmptyObject {
        if value == '{}' {
            return EmptyObject{}
        }
        return error('mcp: expected an empty object payload, got `${value}`')
    } $else $if T is Null {
        if value == null.str() {
            return null
        }
        return error('mcp: expected null, got `${value}`')
    } $else $if T is string {
        if value.len >= 2 && value[0] == `"` && value[value.len - 1] == `"` {
            return json.decode(string, value) or { return err }
        }
        return error('mcp: could not decode `${value}` into string')
    } $else $if T is bool {
        if value == 'true' {
            return true
        }
        if value == 'false' {
            return false
        }
        return error('mcp: could not decode `${value}` into bool')
    } $else {
        return json.decode(T, value) or { return err }
    }
}

fn normalize_client_info(client_info Implementation) Implementation {
    return if client_info.name == '' {
        Implementation{
            name:    default_client_name
            version: if client_info.version == '' {
                default_client_version
            } else {
                client_info.version
            }
        }
    } else if client_info.version == '' {
        Implementation{
            name:    client_info.name
            version: default_client_version
        }
    } else {
        client_info
    }
}

fn normalize_capabilities(capabilities string) string {
    trimmed := capabilities.trim_space()
    return if trimmed.len == 0 { '{}' } else { trimmed }
}

fn normalize_protocol_version(version string) string {
    trimmed := version.trim_space()
    return if trimmed.len == 0 { protocol_version } else { trimmed }
}