utils.py

import click
import os
import json
from functools import wraps
import websocket
import yaml
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse


def print_logo() -> None:
    click.secho(
        r"""      _             _                                _ _ 
     | |           | |                              | (_)
  ___| |_ __ _  ___| | _____ _   _ _ __   ___    ___| |_ 
 / __| __/ _` |/ __| |/ / __| | | | '_ \ / __|  / __| | |
 \__ \ || (_| | (__|   <\__ \ |_| | | | | (__  | (__| | |
 |___/\__\__,_|\___|_|\_\___/\__, |_| |_|\___|  \___|_|_|
                              __/ |                      
                             |___/                       """,
        fg="cyan",
    )


# TODO: Move this somewhere nicer. They're not very sensitive so this will do the trick for now.
DEFAULTS = {
    "local_bridge_url": "https://local-bridge-prod-besg-1011549189584.europe-west1.run.app/",
    "login_url": "https://app.stacksync.com/cdk/login",
    "templates_source_branch": "prod",
}


def get_env_file() -> dict:
    """
    The CLI stores the environment variables in a file called `.stacksync/env.json`.
    This function reads the file and returns the environment variables as a dictionary.
    """
    # If file does not exist, create it with an empty dictionary.
    if not os.path.exists(os.path.join(os.path.expanduser("~/.stacksync"), "env.json")):
        with open(
            os.path.join(os.path.expanduser("~/.stacksync"), "env.json"), "w"
        ) as f:
            json.dump({}, f)
    with open(os.path.join(os.path.expanduser("~/.stacksync"), "env.json"), "r") as f:
        return {**DEFAULTS, **json.load(f)}


def update_env_file(patch: dict) -> None:
    """
    Reads the env file, patches it, and writes it back to the file.
    """
    env_file = get_env_file()
    env_file.update(patch)
    with open(os.path.join(os.path.expanduser("~/.stacksync"), "env.json"), "w") as f:
        json.dump(env_file, f)


def connect_to_local_bridge(
    path: str, *, session_id: str, timeout_s: float = 900.0
) -> websocket.WebSocket:
    """
    Returns a websocket connection to the local bridge.
    """
    base_url = get_env_file()["local_bridge_url"]
    try:
        parsed = urlparse(base_url)
        if parsed.scheme not in ("ws", "wss"):
            raise ValueError("Expected ws:// or wss:// URL")

        # `path` may already include a query string (e.g. "/socket?workspaceId=...")
        # Ensure we always have a leading slash.
        normalized_path = path if path.startswith("/") else f"/{path}"
        path_obj = urlparse(normalized_path)

        query_params = dict(parse_qsl(parsed.query, keep_blank_values=True))
        query_params.update(dict(parse_qsl(path_obj.query, keep_blank_values=True)))
        query_params["sessionId"] = session_id
        query = urlencode(query_params)

        full_url = urlunparse(
            (
                parsed.scheme,
                parsed.netloc,
                path_obj.path or "/",
                "",  # params
                query,
                "",  # fragment
            )
        )
    except Exception as e:
        raise click.ClickException(f"Invalid local bridge URL {base_url!r}: {e}") from e

    try:
        return websocket.create_connection(full_url, timeout=timeout_s)
    except Exception as e:
        raise click.ClickException(
            f"Failed to connect to local bridge at {full_url!r}: {e}"
        ) from e


def respond_to_local_bridge(
    ws: websocket.WebSocket, callback_id: str, payload: dict
) -> None:
    """
    Responds to the local bridge with a payload.
    """
    try:
        ws.send(json.dumps({"callbackId": callback_id, "payload": payload}))
        click.secho("Verifying login...", fg="yellow")
    except Exception as e:
        raise click.ClickException(
            f"Failed to respond to local bridge: {e}"
        ) from e
    return True


def with_auth(func):
    """Inject ``api_key`` from ``~/.stacksync/env.json`` into the decorated callable."""

    @wraps(func)
    def wrapper(*args, **kwargs):
        api_key = get_env_file().get("user_api_key")
        if not api_key:
            raise click.ClickException(
                "Not authenticated. Run `stacksync login` to set your API key."
            )
        return func(*args, **kwargs, api_key=api_key)

    return wrapper


def with_stacksync_yml(func):
    """
    If a `stacksync.yml` file exists in the current directory,
    read it and pass the contents to the decorated function.
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        # Get current directory
        click.echo("Current directory: " + os.getcwd())
        current_dir = os.getcwd()
        stacksync_yml_path = os.path.join(current_dir, "stacksync.yml")
        if not os.path.exists(stacksync_yml_path):
            raise click.ClickException(
                "No stacksync.yml file found. This command must be run in a Stacksync project."
            )
        with open(stacksync_yml_path, "r") as f:
            stacksync_yml = yaml.safe_load(f)
        if not stacksync_yml:
            raise click.ClickException(
                "No `stacksync.yml` file found. Run `stacksync init` to create one."
            )
        return func(*args, **kwargs, stacksync_yml=stacksync_yml)

    return wrapper