app.py

import os
import json
import requests
from fastapi import Request
import websockets
import asyncio
from huggingface_hub import webhook_endpoint

# Speckle stream and authentication information
speckle_token = os.getenv("SPECKLE_TOKEN")
api_url = "https://speckle.xyz/graphql"
stream_id = "1dab2d05eb"
branch_name = "scenario_sycer"

# WebSocket URI
ws_uri = "wss://onlinewebsocketserver.onrender.com"

# Connect to WebSocket and send data
async def send_data(data):
    websocket = None
    try:
        websocket = await websockets.connect(ws_uri)
        print("Connected to the WebSocket server")
        await websocket.send(json.dumps(data))
        print(f"Sent: {data}")
    except Exception as e:
        print(f"Failed to send data: {e}")
    finally:
        if websocket:
            await websocket.close()
            print("WebSocket connection closed.")

def send_graphql_query(speckleToken, apiUrl, query):
    headers = {
        "Authorization": f"Bearer {speckleToken}",
        "Content-Type": "application/json",
        "Accept": "application/json"
    }
    response = requests.post(apiUrl, headers=headers, json={"query": query})
    if response.status_code == 200:
        return response.json()
    else:
        print(f"HTTP Error: {response.status_code}, {response.text}")
        return None

def construct_commit_query(stream_id, branch_name):
    return f"""
    {{
      stream(id: "{stream_id}") {{
        branch(name: "{branch_name}") {{
          commits(limit: 1) {{
            items {{
              id
              message
              referencedObject
            }}
          }}
        }}
      }}
    }}
    """

def fetch_referenced_object(speckleToken, apiUrl, stream_id, object_id):
    query = f"""
    {{
      stream(id: "{stream_id}") {{
        object(id: "{object_id}") {{
          id
          data
        }}
      }}
    }}
    """
    return send_graphql_query(speckleToken, apiUrl, query)

def extract_team_name(object_data):
    data_field = object_data['data']['stream']['object']['data']

    # Check for the first structure with direct 'data'
    try:
        nested_data_str = data_field.get('data')
        if nested_data_str:
            nested_data = json.loads(nested_data_str)
            team_name = nested_data.get('teamName')
            if team_name:
                return team_name
    except (KeyError, json.JSONDecodeError):
        pass

    # Check for the second structure under 'Data' -> '@{0}'
    try:
        data_object = data_field.get('Data', {})
        team_name_list = data_object.get('@{0}')
        if team_name_list and isinstance(team_name_list, list):
            return team_name_list[0]
    except (KeyError, json.JSONDecodeError) as e:
        print(f"Error extracting team name: {e}")

    return None

@webhook_endpoint
async def update_streams(request: Request):
    payload = await request.json()
    print(f"Received webhook payload: {payload}")

    commit_query = construct_commit_query(stream_id, branch_name)
    result = send_graphql_query(speckle_token, api_url, commit_query)

    if result and 'data' in result:
        commit_data = result['data']['stream']['branch']['commits']['items'][0]
        referenced_object_id = commit_data['referencedObject']
        print(f"Referenced Object ID: {referenced_object_id}")

        object_data = fetch_referenced_object(speckle_token, api_url, stream_id, referenced_object_id)
        if object_data:
            print(json.dumps(object_data, indent=2))
            
            team_name = extract_team_name(object_data)
            
            if team_name:
                print(f"SENDING Team Name: {team_name}")
                await send_data(team_name)
            else:
                print("Team name not found in the object data.")
        else:
            print("Failed to retrieve the referenced object data.")
    else:
        print("Failed to retrieve commit data.")

    return "Data sent successfully!"

# Uncomment below if you want to use Gradio for a manual interface
"""
import gradio as gr
iface = gr.Interface(
    fn=update_streams,
    inputs=gr.components.Button(value="Update Streams"),
    outputs="text",
)
iface.launch()
"""