"""
Manual T4 verification: call the RAG tool THROUGH the AgentCore Gateway (MCP) and
print the de-identified evidence package. This is the "manual agentcore invoke"
the task's acceptance asks for -- it exercises the real Gateway -> Lambda path,
not the Lambda directly.

Auth is AWS_IAM: every MCP request is SigV4-signed for the `bedrock-agentcore`
service with your local creds. MCP uses streamable HTTP (JSON-RPC, responses may
come back as SSE), so we do initialize -> tools/list -> tools/call.

Usage:
    REGION=ap-southeast-1 AGENTCORE_GATEWAY_URL=<url> \
        python3 scripts/gateway_invoke_test.py CUST_000123

Requires botocore (pip install boto3). Depends on nothing else.
"""
import json
import os
import sys
import urllib.request

from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.session import Session

REGION = os.environ["REGION"]
URL = os.environ["AGENTCORE_GATEWAY_URL"]
TOKEN = sys.argv[1] if len(sys.argv) > 1 else "CUST_000123"

_creds = Session().get_credentials()
_session_id = None  # set from the initialize response


def _rpc(method, params=None, notify=False):
    """One SigV4-signed MCP JSON-RPC call. Returns the parsed result (or None)."""
    global _session_id
    body = {"jsonrpc": "2.0", "method": method}
    if not notify:
        body["id"] = 1
    if params is not None:
        body["params"] = params
    data = json.dumps(body).encode()

    headers = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"}
    if _session_id:
        headers["Mcp-Session-Id"] = _session_id

    req = AWSRequest(method="POST", url=URL, data=data, headers=headers)
    SigV4Auth(_creds, "bedrock-agentcore", REGION).add_auth(req)

    try:
        resp = urllib.request.urlopen(urllib.request.Request(URL, data=data, headers=dict(req.headers)), timeout=40)
    except urllib.error.HTTPError as e:
        print(f"HTTP {e.code} on {method}: {e.read().decode()[:400]}", file=sys.stderr)
        raise

    sid = resp.headers.get("Mcp-Session-Id")
    if sid:
        _session_id = sid
    raw = resp.read().decode()
    if notify:
        return None

    # Response is either JSON or SSE (event-stream). Extract the JSON-RPC object.
    payload = raw
    if "text/event-stream" in (resp.headers.get("Content-Type") or "") or raw.lstrip().startswith("event:"):
        for line in raw.splitlines():
            if line.startswith("data:"):
                payload = line[5:].strip()
    obj = json.loads(payload)
    if "error" in obj:
        raise RuntimeError(f"MCP error on {method}: {obj['error']}")
    return obj.get("result")


def main():
    _rpc("initialize", {
        "protocolVersion": "2025-06-18",
        "capabilities": {},
        "clientInfo": {"name": "hncb-t4-test", "version": "1.0"},
    })
    _rpc("notifications/initialized", notify=True)

    tools = _rpc("tools/list").get("tools", [])
    names = [t["name"] for t in tools]
    print("tools exposed by gateway:", names)
    # Gateway namespaces tool names (e.g. rag___get_customer_activity_summary).
    tool = next((n for n in names if n.endswith("get_customer_activity_summary")), None)
    if not tool:
        sys.exit("get_customer_activity_summary not found on the gateway")

    result = _rpc("tools/call", {"name": tool, "arguments": {"customer_token": TOKEN}})
    print(f"\n=== evidence package for {TOKEN} (via Gateway MCP) ===")
    # MCP tool results come back as content blocks; print text blocks.
    for block in result.get("content", []):
        if block.get("type") == "text":
            try:
                print(json.dumps(json.loads(block["text"]), indent=2, ensure_ascii=False))
            except json.JSONDecodeError:
                print(block["text"])


if __name__ == "__main__":
    main()