#!/usr/bin/env python3
"""Complete read-only Upkeep OAuth PKCE + MCP example (Python 3.9+, no dependencies).

Run: python3 upkeep-connect.py --query fridge
This registers a public read-only client, binds a loopback listener, opens the
Upkeep sign-in/consent flow, exchanges the resource-bound code, initializes MCP,
then discovers tools and searches assets. Tokens stay in memory and are never
printed or written. The example never writes inventory or buys membership.
"""
import argparse
import getpass
import json
import os
import sys
from urllib import error, parse, request

ENDPOINT = "https://apps.sharshi.com/api/v1/upkeep/mcp"
MAX_RESPONSE_BYTES = 3 * 1024 * 1024


class McpError(Exception):
    """A failed transport/protocol/tool response, never an empty inventory."""


def decode_response(status, content_type, body, request_id):
    if not 200 <= status < 300:
        explanations = {401: "Reconnect: the credential is invalid or expired.",
                        402: "Active membership is required.",
                        403: "The credential lacks the required scope; check WWW-Authenticate.",
                        406: "Send Accept: application/json, text/event-stream."}
        raise McpError(f"HTTP {status}. {explanations.get(status, 'Request failed.')}")
    if request_id is None and status in (202, 204) and not body.strip():
        return None  # Accepted initialized notification, not an inventory result.
    try:
        if content_type.split(";", 1)[0].strip() == "text/event-stream":
            messages, data = [], []
            # SSE may contain comments, notifications, and multiline data fields.
            for line in body.replace("\r\n", "\n").split("\n") + [""]:
                if not line:
                    if data:
                        messages.append(json.loads("\n".join(data)))
                        data = []
                elif line.startswith("data:"):
                    data.append(line[5:].removeprefix(" "))
        elif content_type.split(";", 1)[0].strip() == "application/json":
            messages = [json.loads(body)]
        else:
            raise McpError("Expected JSON or SSE, not an HTML page or another content type.")
    except (ValueError, TypeError) as exc:
        raise McpError("Unreadable JSON/SSE response; no inventory result was obtained.") from exc

    for message in messages:
        if not isinstance(message, dict) or message.get("jsonrpc") != "2.0":
            raise McpError("Invalid JSON-RPC response.")
        if message.get("error") is not None:
            rpc_error = message["error"]
            code = rpc_error.get("code") if isinstance(rpc_error, dict) else "unknown"
            raise McpError(f"JSON-RPC error {code}; check the method, discovered tool and input schema.")
        if message.get("id") != request_id:
            continue
        result = message.get("result")
        if not isinstance(result, dict):
            raise McpError("Missing JSON-RPC result; do not substitute an empty list.")
        if result.get("isError") is True:
            raise McpError("Tool execution failed (isError=true); no successful inventory result.")
        return result
    raise McpError("No completed response with the matching request ID.")


def tool_data(result):
    if "structuredContent" in result:
        value = result["structuredContent"]
        if not isinstance(value, dict):
            raise McpError("Invalid structuredContent; expected a JSON object.")
        return value
    # Older clients/servers expose the same object as serialized JSON text.
    values = []
    contents = result.get("content", [])
    if not isinstance(contents, list):
        raise McpError("Invalid text content; expected an array.")
    for content in contents:
        if isinstance(content, dict) and content.get("type") == "text":
            try:
                value = json.loads(content.get("text", ""))
            except (ValueError, TypeError):
                continue
            if isinstance(value, dict):
                values.append(value)
    if len(values) != 1:
        raise McpError("Expected one JSON text object; no valid tool data was obtained.")
    return values[0]


def search_inventory(client, tools, query, limit=25):
    legacy = next((tool for tool in tools if tool.get("name") == "upkeep.search_assets"), None)
    if legacy:
        return tool_data(client.call("tools/call", {"name": legacy["name"], "arguments": {"query": query, "limit": limit}}))
    workflow = next((tool for tool in tools if tool.get("name") == "morrowfield.assets"), None)
    if not workflow:
        raise McpError("Inventory workflow is absent; no action name will be guessed.")
    discovered = tool_data(client.call("tools/call", {"name": workflow["name"], "arguments": {}}))
    operation = next((action for action in discovered.get("actions", []) if action.get("operation") == "upkeep.search_assets"), None)
    if not operation:
        raise McpError("Inventory search is not available in this connection.")
    data = tool_data(client.call("tools/call", {"name": workflow["name"], "arguments": {"operation": operation["operation"], "arguments": {"query": query, "limit": limit}}}))
    if not isinstance(data.get("result"), dict):
        raise McpError("Inventory workflow did not return a completed result.")
    return data["result"]


def asset_rows(data):
    rows = data.get("items")
    if not isinstance(rows, list) or any(not isinstance(row, dict) for row in rows):
        raise McpError("Search result has no valid items array; this is not empty inventory.")
    if any(not isinstance(row.get("id"), str) or not isinstance(row.get("name"), str) for row in rows):
        raise McpError("Search rows must contain opaque string IDs and names.")
    if type(data.get("count")) is not int or data["count"] != len(rows):
        raise McpError("Search result count does not match the returned items.")
    return rows


class NoRedirect(request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        # Never forward the Bearer credential to a redirect destination.
        raise McpError("Endpoint redirected; verify the canonical server URL before reconnecting.")


class Client:
    def __init__(self, endpoint, token):
        url = parse.urlsplit(endpoint)
        if url.username or url.password or url.query or url.fragment or not url.hostname:
            raise McpError("Use a server URL without credentials, query parameters or fragments.")
        if url.scheme != "https" and not (url.scheme == "http" and url.hostname in ("localhost", "127.0.0.1", "::1")):
            raise McpError("HTTPS is required outside loopback development.")
        self.endpoint, self.token = endpoint, token
        self.protocol_version, self.next_id = None, 1
        self.opener = request.build_opener(NoRedirect())

    def call(self, method, params, notification=False):
        request_id = None if notification else self.next_id
        if not notification:
            self.next_id += 1
        payload = {"jsonrpc": "2.0", "method": method, "params": params}
        if request_id is not None:
            payload["id"] = request_id  # Integers also avoid this SDK's string-ID restriction.
        headers = {"Authorization": f"Bearer {self.token}", "Content-Type": "application/json",
                   "Accept": "application/json, text/event-stream", "X-Agent-Name": "upkeep-read-diagnostic"}
        if self.protocol_version:
            headers["MCP-Protocol-Version"] = self.protocol_version
        req = request.Request(self.endpoint, data=json.dumps(payload).encode(), headers=headers, method="POST")
        try:
            with self.opener.open(req, timeout=30) as response:
                raw = response.read(MAX_RESPONSE_BYTES + 1)
                if len(raw) > MAX_RESPONSE_BYTES:
                    raise McpError("Response exceeds the diagnostic's size limit; use a narrower query.")
                return decode_response(response.status, response.headers.get("Content-Type", ""), raw.decode("utf-8"), request_id)
        except error.HTTPError as exc:
            # Report the status, never echo a body or token into diagnostics/logs.
            return decode_response(exc.code, "", "", request_id)
        except (error.URLError, TimeoutError, UnicodeError) as exc:
            raise McpError("Connection or decoding failed; this says nothing about saved inventory.") from exc

    def connect(self):
        result = self.call("initialize", {"protocolVersion": "2025-06-18", "capabilities": {},
                           "clientInfo": {"name": "upkeep-read-diagnostic", "version": "1.0"}})
        self.protocol_version = result.get("protocolVersion")
        if not isinstance(self.protocol_version, str):
            raise McpError("Initialization did not negotiate a protocol version.")
        self.call("notifications/initialized", {}, notification=True)
        tools, cursor = [], None
        for _ in range(10):
            page = self.call("tools/list", {"cursor": cursor} if cursor else {})
            if not isinstance(page.get("tools"), list) or any(not isinstance(tool, dict) for tool in page["tools"]):
                raise McpError("Invalid tool catalog.")
            tools.extend(page["tools"])
            cursor = page.get("nextCursor")
            if not cursor:
                return tools
        raise McpError("Tool catalog pagination limit reached.")


import base64
import hashlib
import secrets
import time
import webbrowser
from http.server import HTTPServer, BaseHTTPRequestHandler


def validate_base(value):
    url = parse.urlsplit(value)
    if url.username or url.password or url.query or url.fragment or url.path not in ("", "/") or not url.hostname:
        raise McpError("Use an origin URL without credentials or a path.")
    if url.scheme != "https" and not (url.scheme == "http" and url.hostname in ("127.0.0.1", "localhost", "::1")):
        raise McpError("HTTPS is required outside loopback development.")
    return value.rstrip("/")


def endpoint_on_origin(url, base):
    parsed, origin = parse.urlsplit(url), parse.urlsplit(base)
    if (parsed.scheme, parsed.netloc) != (origin.scheme, origin.netloc) or parsed.username or parsed.password or parsed.fragment or parsed.query:
        raise McpError("Discovery returned an endpoint outside the configured origin.")
    return url


def json_request(opener, url, payload=None, form=False):
    headers = {"Accept": "application/json"}
    data = None
    if payload is not None:
        headers["Content-Type"] = "application/x-www-form-urlencoded" if form else "application/json"
        data = parse.urlencode(payload).encode() if form else json.dumps(payload).encode()
    req = request.Request(url, data=data, headers=headers)
    try:
        with opener.open(req, timeout=30) as response:
            if response.headers.get_content_type() != "application/json":
                raise McpError("OAuth endpoint did not return JSON.")
            raw = response.read(MAX_RESPONSE_BYTES + 1)
            if len(raw) > MAX_RESPONSE_BYTES:
                raise McpError("OAuth response exceeded its size limit.")
            result = json.loads(raw)
            if not isinstance(result, dict) or result.get("error"):
                raise McpError("OAuth operation failed; restart connection.")
            return result
    except error.HTTPError as exc:
        raise McpError(f"OAuth HTTP {exc.code}; check configuration or restart connection.") from None
    except (error.URLError, TimeoutError, ValueError):
        raise McpError("OAuth connection or response decoding failed.") from None


def callback_handler(expected_state):
    class Callback(BaseHTTPRequestHandler):
        def log_message(self, *args):
            pass  # Never log a callback URL containing an authorization code.

        def do_GET(self):
            url = parse.urlsplit(self.path)
            values = parse.parse_qs(url.query)
            valid = url.path == "/callback" and len(values.get("state", [])) == 1 and secrets.compare_digest(values["state"][0], expected_state)
            code = values.get("code", [])
            if not valid or len(code) != 1 or not code[0] or "error" in values:
                self.send_response(400)
                body = b"Connection rejected. Return to Upkeep and restart authorization."
            else:
                self.server.authorization_code = code[0]
                self.send_response(200)
                body = b"Connection approved. Return to the terminal; this window can be closed."
            self.send_header("Content-Type", "text/plain; charset=utf-8")
            self.send_header("Cache-Control", "no-store")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
    return Callback


def oauth_connect(base_url, browser_open=webbrowser.open, timeout=300, announce=print):
    base = validate_base(base_url)
    opener = request.build_opener(NoRedirect())
    state = secrets.token_urlsafe(32)
    verifier = secrets.token_urlsafe(64)
    challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
    resource = base + "/api/v1/upkeep/mcp"
    # Bind before registration/authorization, so callbacks cannot race startup.
    with HTTPServer(("127.0.0.1", 0), callback_handler(state)) as listener:
        listener.timeout = 1
        listener.authorization_code = None
        redirect_uri = f"http://127.0.0.1:{listener.server_port}/callback"
        metadata = json_request(opener, base + "/.well-known/oauth-authorization-server")
        registration = endpoint_on_origin(metadata["registration_endpoint"], base)
        token_endpoint = endpoint_on_origin(metadata["token_endpoint"], base)
        client = json_request(opener, registration, {"client_name": "Upkeep read-only example", "redirect_uris": [redirect_uri], "scope": "upkeep:read"})
        client_id = client.get("client_id")
        if not isinstance(client_id, str) or not client_id:
            raise McpError("Registration returned no client ID.")
        announce("1. Registered a read-only public client (HTTP 201, client_id).")
        parameters = {"client_id": client_id, "redirect_uri": redirect_uri, "response_type": "code", "scope": "upkeep:read", "state": state, "code_challenge": challenge, "code_challenge_method": "S256", "resource": resource}
        authorization_url = base + "/upkeep-ai/start?" + parse.urlencode(parameters)
        announce("2. Complete email sign-in and explicit consent in Upkeep. If membership is required, the user decides whether to purchase; this script never purchases.")
        if not browser_open(authorization_url):
            announce("Open this Upkeep authorization URL in your local browser:")
            announce(authorization_url)  # Contains no token, code or PKCE verifier.
        deadline = time.monotonic() + timeout
        while listener.authorization_code is None and time.monotonic() < deadline:
            listener.handle_request()
        if listener.authorization_code is None:
            raise McpError("Consent timed out. Restart the example; no token was obtained.")
        tokens = json_request(opener, token_endpoint, {"grant_type": "authorization_code", "client_id": client_id, "code": listener.authorization_code, "redirect_uri": redirect_uri, "code_verifier": verifier, "resource": resource}, form=True)
        token = tokens.get("access_token")
        if not isinstance(token, str) or not token or tokens.get("token_type", "").lower() != "bearer":
            raise McpError("Exchange did not return a bearer credential.")
        announce("3. State verified; exchanged PKCE code (HTTP 200, bearer credential held only in memory).")
        return token


def connect_and_read(base, query="", browser_open=webbrowser.open, timeout=300, announce=print):
    token = oauth_connect(base, browser_open=browser_open, timeout=timeout, announce=announce)
    client = Client(validate_base(base) + "/api/v2/upkeep/mcp", token)
    tools = client.connect()
    announce(f"4. initialize negotiated {client.protocol_version}; notifications/initialized accepted; tools/list returned {len(tools)} tools.")
    data = search_inventory(client, tools, query)
    rows = asset_rows(data)
    announce(f"5. tools/call succeeded; structuredContent.items contains {len(rows)} assets in this page.")
    if data.get("search", {}).get("status") == "pending":
        announce("Ranking pending: repeat the query before concluding no matches exist.")
    if data.get("meta", {}).get("next_cursor"):
        announce("More records exist; follow structuredContent.meta.next_cursor.")
    return rows


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default="https://apps.sharshi.com", help="HTTPS Upkeep origin; loopback HTTP allowed for development")
    parser.add_argument("--query", default="")
    parser.add_argument("--no-browser", action="store_true", help="Print the Upkeep authorization URL instead of opening it")
    args = parser.parse_args()
    try:
        rows = connect_and_read(args.base_url, args.query, browser_open=(lambda _: False) if args.no_browser else webbrowser.open)
        for row in rows:
            print(json.dumps({"id": row["id"], "name": row["name"]}, ensure_ascii=True))
        print("Read-only example complete. Disconnect the example in Connected applications when finished.")
        return 0
    except (McpError, KeyError, OSError):
        print("Connection failed. Verify Upkeep configuration or restart sign-in/consent. No inventory success is claimed.", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())