#!/usr/bin/env python3
import argparse
import json
import sys
import urllib.request
import urllib.error

DEFAULT_PROTOCOL_VERSIONS = ["2025-03-26", "2024-11-05"]

SSE_ACCEPT = "application/json, text/event-stream"
CONTENT_TYPE = "application/json"


def sse_read_until_request_id(resp, want_id, timeout_seconds=10):
    found = None
    buf_data_lines = []

    # SSE: event/data lines, with blank line separating messages.
    while True:
        line = resp.readline()
        if not line:
            break  # EOF

        s = line.decode("utf-8", errors="replace").rstrip("\r\n")

        if s == "":
            if buf_data_lines:
                data_text = "\n".join(buf_data_lines)
                buf_data_lines = []
                try:
                    obj = json.loads(data_text)
                    _id = obj.get("id")
                    if _id == want_id or str(_id) == str(want_id):
                        if "result" in obj or "error" in obj:
                            found = obj
                            break
                except json.JSONDecodeError:
                    pass
            continue

        if s.startswith("data:"):
            buf_data_lines.append(s[len("data:"):].lstrip())
        # ignore "event:" and other fields

    # In case server ends without a trailing blank line, parse last buffered message
    if found is None and buf_data_lines:
        data_text = "\n".join(buf_data_lines)
        try:
            obj = json.loads(data_text)
            _id = obj.get("id")
            if _id == want_id or str(_id) == str(want_id):
                if "result" in obj or "error" in obj:
                    found = obj
        except json.JSONDecodeError:
            pass

    return found


def post_and_read_sse(url, body_obj, headers, want_id, timeout=10):
    data = json.dumps(body_obj).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")

    with urllib.request.urlopen(req, timeout=timeout) as resp:
        session_id = resp.headers.get("mcp-session-id") or resp.headers.get("Mcp-Session-Id")
        msg = sse_read_until_request_id(resp, want_id=want_id, timeout_seconds=timeout)
        return session_id, msg


def post_json_no_read(url, body_obj, headers, timeout=10):
    data = json.dumps(body_obj).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        # Consume at least headers; ignore body
        _ = resp.status


def delete_no_read(url, headers, timeout=10):
    req = urllib.request.Request(url, method="DELETE", headers=headers)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.status


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("endpoint", help="Full MCP streamable HTTP endpoint URL, e.g. http://host:8010/mcp")
    ap.add_argument("--tools", action="store_true", help="Also call tools/list after initialize")
    args = ap.parse_args()

    endpoint = args.endpoint
    base_headers = {
        "Content-Type": CONTENT_TYPE,
        "Accept": SSE_ACCEPT,
        "User-Agent": "test-mcp/1.0",
    }

    init_id = 1
    tools_id = 2
    init_result = None
    session_id = None

    # ---- initialize (POST) ----
    last_err = None
    for protocol_version in DEFAULT_PROTOCOL_VERSIONS:
        init_body = {
            "jsonrpc": "2.0",
            "id": init_id,
            "method": "initialize",
            "params": {
                "protocolVersion": protocol_version,
                "capabilities": {},
                "clientInfo": {"name": "test-mcp", "version": "1.0.0"},
            },
        }
        try:
            session_id, init_result = post_and_read_sse(endpoint, init_body, base_headers, want_id=init_id, timeout=10)
            if init_result is not None:
                break
        except urllib.error.HTTPError as e:
            last_err = e
        except Exception as e:
            last_err = e

    if init_result is None:
        if last_err:
            print(f"Initialize failed: {last_err}")
        else:
            print("Initialize failed: no initialize response received")
        sys.exit(1)

    print("Initialize response:")
    print(json.dumps(init_result, indent=2, sort_keys=True))

    if not session_id:
        print("Warning: MCP session id header not found; termination and follow-up calls may fail.")
        session_id = None

    # ---- initialized (notification, POST) ----
    if session_id:
        headers = dict(base_headers)
        headers["Mcp-Session-Id"] = session_id
        initialized_body = {"jsonrpc": "2.0", "method": "initialized", "params": {}}
        post_json_no_read(endpoint, initialized_body, headers, timeout=10)

    # ---- optional tools/list ----
    if args.tools and session_id:
        headers = dict(base_headers)
        headers["Mcp-Session-Id"] = session_id
        tools_body = {"jsonrpc": "2.0", "id": tools_id, "method": "tools/list", "params": {}}
        _, tools_result = post_and_read_sse(endpoint, tools_body, headers, want_id=tools_id, timeout=10)
        print("tools/list response:")
        if tools_result is None:
            print("(no tools/list SSE response received)")
        else:
            print(json.dumps(tools_result, indent=2, sort_keys=True))

    # ---- terminate session (DELETE) ----
    if session_id:
        term_headers = {
            "Accept": "application/json",
            "User-Agent": "test-mcp/1.0",
            "Mcp-Session-Id": session_id,
        }
        try:
            status = delete_no_read(endpoint, term_headers, timeout=10)
            print(f"Session termination: HTTP {status}")
        except urllib.error.HTTPError as e:
            print(f"Session termination failed: HTTP {e.code}")
            try:
                body = e.read().decode("utf-8", errors="replace")
                if body:
                    print(body)
            except Exception:
                pass


if __name__ == "__main__":
    main()