main.py

# /// script
# dependencies = [
#   "azure-messaging-webpubsubclient==1.1.0",
#   "protobuf==6.33.1",
#   "requests==2.32.5",
#   "zstandard==0.25.0",
# ]
# ///

import json
import os
import threading
import time
from pathlib import Path
from typing import Dict, List, Optional

import requests
from azure.messaging.webpubsubclient import WebPubSubClient
from azure.messaging.webpubsubclient.models import (
    CallbackType,
    OnConnectedArgs,
    OnDisconnectedArgs,
    OnGroupDataMessageArgs,
)
from google.protobuf import any_pb2

from decompression_utils import (
    decompress_gex_message,
    decompress_greek_message,
    decompress_orderflow_message,
    decompress_spot_message,
)
from group_config import (
    add_colocated_spot_groups,
    build_patch_groups,
    build_post_groups,
    count_group_memberships,
)

# --- Configuration ---

SCRIPT_DIR = Path(__file__).resolve().parent


def load_dotenv_file(path: Path) -> None:
    """Load local environment variables without requiring an extra dependency."""
    if not path.exists():
        return

    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue

        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key:
            os.environ.setdefault(key, value)


load_dotenv_file(SCRIPT_DIR / ".env")


def required_environment_variable(name: str) -> str:
    """Return one required environment variable or raise a clear error."""
    value = os.environ.get(name)
    if value is None or not value.strip():
        raise RuntimeError(f"{name} is required and must not be empty.")
    return value.strip()


API_KEY = required_environment_variable("GEXBOT_API_KEY")
USER_AGENT = required_environment_variable("GEXBOT_USER_AGENT")
BASE_URL = required_environment_variable("BASE_URL").rstrip("/")
NEGOTIATE_URL = f"{BASE_URL}/negotiate"

# Configure analytics only. The client adds one matching {ticker}_spot group to
# each V2 hub so analytics and spot use separate groups on the same connection.
INITIAL_ANALYTICS_GROUP_CONFIG: Dict[str, List[str]] = {
    "v2_classic": [],
    "v2_state_gex": [
        "SPX_state_gex_full",
        "SPX_state_gex_zero",
    ],
    "v2_state_greeks_zero": [
        "SPX_state_gamma_zero",
    ],
    "v2_state_greeks_one": [],
    "v2_orderflow": [],
}

# PATCH /negotiate is a full replacement. The generated payload includes the
# complete analytics and colocated spot membership set.
UPDATED_ANALYTICS_GROUP_CONFIG: Dict[str, List[str]] = {
    "v2_classic": [
        "SPX_classic_gex_full",
    ],
    "v2_state_gex": [
        "NDX_state_gex_full",
    ],
    "v2_state_greeks_zero": [
        "SPX_state_delta_zero",
        "SPX_state_gamma_zero",
    ],
    "v2_state_greeks_one": [],
    "v2_orderflow": [],
}

# Set to 0 to skip the PATCH replacement example.
PATCH_UPDATE_AFTER_SECONDS = 20


def print_group_config(title: str, group_config: Dict[str, List[str]]) -> None:
    """Print the final analytics and spot memberships for each V2 hub."""
    print(f"\n--- {title} ---")
    for hub_key, groups in group_config.items():
        if groups:
            print(f"[{hub_key}] {', '.join(groups)}")
    print(f"Total hub memberships: {count_group_memberships(group_config)}")


def build_auth_headers(api_key: str, user_agent: str) -> Dict[str, str]:
    """Build the required API authorization and content-negotiation headers."""
    if not api_key or api_key == "your_api_key_here":
        raise ValueError("GEXBOT_API_KEY is not set.")

    if not user_agent.strip():
        raise ValueError('GEXBOT_USER_AGENT is not set. Example: "AcmeQuantClient/1.0"')

    return {
        "Authorization": f"Bearer {api_key}",
        "User-Agent": user_agent.strip(),
        "Accept": "application/json",
    }


def post_negotiate(
    api_key: str, user_agent: str, group_config: Dict[str, List[str]]
) -> Dict:
    """Negotiate V2 connections with analytics and one spot request per ticker."""
    groups = build_post_groups(group_config)
    if not groups:
        raise ValueError("No initial websocket groups configured.")

    headers = build_auth_headers(api_key, user_agent)
    payload = {"groups": groups}

    print(f"POST {NEGOTIATE_URL} with {len(groups)} group(s)...")
    response = requests.post(NEGOTIATE_URL, headers=headers, json=payload, timeout=30)
    response.raise_for_status()
    return response.json()


def patch_replace_groups(
    api_key: str, user_agent: str, group_config: Dict[str, List[str]]
) -> Dict:
    """Replace all V2 analytics and colocated spot memberships."""
    groups = build_patch_groups(group_config)
    headers = build_auth_headers(api_key, user_agent)
    payload = {"groups": groups}

    print(f"PATCH {NEGOTIATE_URL} with {len(groups)} group(s)...")
    response = requests.patch(NEGOTIATE_URL, headers=headers, json=payload, timeout=30)
    response.raise_for_status()
    return response.json()


class WebPubSubClientManager:
    """Manages one Azure Web PubSub hub connection."""

    def __init__(self, hub_key: str, connection_url: str, subscribed_groups: List[str]):
        """Create one manager for a signed V2 hub connection URL."""
        self.hub_key = hub_key
        self.subscribed_groups = subscribed_groups
        self.client = WebPubSubClient(connection_url)
        self.thread: Optional[threading.Thread] = None

        self.client.subscribe(CallbackType.CONNECTED, self.on_connected)
        self.client.subscribe(CallbackType.DISCONNECTED, self.on_disconnected)
        self.client.subscribe(CallbackType.GROUP_MESSAGE, self.on_group_message)

    def start(self):
        """Open the managed connection on a background thread."""
        print(f"[{self.hub_key}] Starting client...")
        self.thread = threading.Thread(target=self.client.open, daemon=True)
        self.thread.start()

    def stop(self):
        """Close the managed connection."""
        print(f"[{self.hub_key}] Stopping client...")
        self.client.close()

    def update_subscribed_groups(self, subscribed_groups: List[str]):
        """Update the local membership labels after a successful PATCH request."""
        self.subscribed_groups = subscribed_groups
        active_groups = ", ".join(subscribed_groups) or "none"
        print(f"[{self.hub_key}] Active groups now: {active_groups}")

    def on_connected(self, event: OnConnectedArgs):
        """Report a successful Web PubSub connection."""
        print(f"[{self.hub_key}] ✅ connected (ID: {event.connection_id})")
        if self.subscribed_groups:
            print(f"[{self.hub_key}] subscribed: {', '.join(self.subscribed_groups)}")

    def on_disconnected(self, event: OnDisconnectedArgs):
        """Report a closed Web PubSub connection."""
        print(f"[{self.hub_key}] ❌ disconnected: {event.message}")

    def on_group_message(self, event: OnGroupDataMessageArgs):
        """Decode and print one analytics or spot group message."""
        try:
            any_message = any_pb2.Any()
            any_message.ParseFromString(event.data)
            message_type_url = any_message.type_url

            current_category = ""
            for package_name in ["classic", "state", "orderflow"]:
                separator = f"_{package_name}_"
                if separator in event.group:
                    current_category = event.group.split(separator)[-1]
                    break

            if "proto.spot" in message_type_url:
                spot_data = decompress_spot_message(any_message)
                print(
                    f"[{self.hub_key}] SPOT: {spot_data.get('ticker')} "
                    f"@ {spot_data.get('spot')} timestamp={spot_data.get('timestamp')}"
                )
                return

            if "proto.gex" in message_type_url:
                gex_data = decompress_gex_message(any_message)
                if gex_data:
                    ticker = gex_data.get("ticker")
                    timestamp = gex_data.get("timestamp")
                    print(f"[{self.hub_key}] GEX: {ticker} timestamp={timestamp}")
                return

            if "proto.greek" in message_type_url:
                greek_data = decompress_greek_message(any_message, current_category)
                if greek_data:
                    ticker = greek_data.get("ticker")
                    print(f"[{self.hub_key}] {current_category}: {ticker}")
                return

            if "proto.orderflow" in message_type_url:
                orderflow_data = decompress_orderflow_message(any_message)
                if orderflow_data:
                    print(
                        f"[{self.hub_key}] Orderflow: "
                        f"{orderflow_data.get('ticker')} "
                        f"timestamp={orderflow_data.get('timestamp')}"
                    )
                return

            print(f"[{self.hub_key}] Unknown message type_url: {message_type_url}")
        except Exception as exc:
            print(f"[{self.hub_key}] Failed to parse message: {exc}")


def start_clients(
    websocket_urls: Dict[str, str], group_config: Dict[str, List[str]]
) -> Dict[str, WebPubSubClientManager]:
    """Open one managed client for each returned V2 hub URL."""
    managers: Dict[str, WebPubSubClientManager] = {}

    for hub_key, url in websocket_urls.items():
        manager = WebPubSubClientManager(
            hub_key,
            url,
            group_config.get(hub_key, []),
        )
        manager.start()
        managers[hub_key] = manager

    return managers


def update_manager_labels(
    managers: Dict[str, WebPubSubClientManager], group_config: Dict[str, List[str]]
):
    """Apply successful PATCH membership labels to all managed clients."""
    for hub_key, manager in managers.items():
        manager.update_subscribed_groups(group_config.get(hub_key, []))


if __name__ == "__main__":
    client_managers: Dict[str, WebPubSubClientManager] = {}

    try:
        initial_group_config = add_colocated_spot_groups(INITIAL_ANALYTICS_GROUP_CONFIG)
        updated_group_config = add_colocated_spot_groups(UPDATED_ANALYTICS_GROUP_CONFIG)

        print_group_config("Initial POST memberships", initial_group_config)
        negotiate_data = post_negotiate(API_KEY, USER_AGENT, initial_group_config)

        websocket_urls = negotiate_data.get("websocket_urls")
        if not websocket_urls:
            raise RuntimeError("No websocket_urls returned from negotiate.")

        print("\n--- Starting WebSocket clients ---")
        client_managers = start_clients(websocket_urls, initial_group_config)

        if PATCH_UPDATE_AFTER_SECONDS > 0:
            print(
                "\nWaiting "
                f"{PATCH_UPDATE_AFTER_SECONDS}s before the PATCH replacement example..."
            )
            time.sleep(PATCH_UPDATE_AFTER_SECONDS)

            print_group_config("PATCH replacement memberships", updated_group_config)
            update_result = patch_replace_groups(
                API_KEY, USER_AGENT, updated_group_config
            )
            print("PATCH replacement result:")
            print(json.dumps(update_result, indent=2))
            update_manager_labels(client_managers, updated_group_config)

        print("\nClients are running. Press Ctrl+C to stop.")
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        pass
    except requests.exceptions.HTTPError as http_err:
        print(f"HTTP error: {http_err}")
        if http_err.response is not None:
            print(http_err.response.text)
    except Exception as exc:
        print(f"Error: {exc}")
    finally:
        if client_managers:
            print("\n--- Stopping Clients ---")
            for manager in client_managers.values():
                manager.stop()
            print("All clients stopped.")