#!/usr/bin/env python3
# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

"""
Standalone SSH ProxyCommand helper.

Reads the JWT cookie from the datatailr remote CLI config and exec's
websocat to tunnel SSH traffic over a WebSocket connection.

Before launching websocat, checks if the target workstation is running
and starts it if necessary, showing progress on stderr (stdout is
reserved for the SSH data channel).

This script is copied into ~/.dt/remote_client/ by ``datatailr setup-ssh``
and referenced from the SSH config ProxyCommand directive.
"""

from __future__ import annotations

import base64
import json
import os
import pathlib
import random
import shutil
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

CONFIG_PATH = pathlib.Path.home() / ".dt" / "remote_client" / "remote_client.cfg"
COOKIE_NAME = "X-Datatailr-Oidc-Data"

_LOADING_PHRASES = [
    "Warming up the engines",
    "Brewing some fresh compute",
    "Waking up the cloud hamsters",
    "Allocating quantum flux capacitors",
    "Spinning up your workstation",
    "Teaching electrons new tricks",
    "Negotiating with the cloud gods",
    "Compiling the universe... almost there",
    "Reticulating splines",
    "Convincing servers to cooperate",
    "Aligning the bits and bytes",
    "Summoning containers from the void",
    "Deploying caffeinated microservices",
    "Polishing your dev environment",
    "Downloading more RAM... just kidding",
    "Herding digital cats into containers",
    "Untangling the network cables",
    "Charging the laser beams",
    "Tuning the hyperparameters of reality",
    "Asking the cloud nicely",
    "Performing ancient deployment rituals",
    "Consulting the documentation... wait, there is none",
    "Assembling your digital workstation",
    "Provisioning pocket dimensions",
    "Loading awesomeness at maximum bandwidth",
]


def _resolve_websocat() -> str | None:
    found = shutil.which("websocat")
    if found:
        return found

    candidates = [
        pathlib.Path("/usr/local/bin/websocat"),
        pathlib.Path.home() / ".local" / "bin" / "websocat",
        pathlib.Path.home() / "scoop" / "shims" / "websocat.exe",
    ]
    for candidate in candidates:
        if candidate.exists():
            return str(candidate)
    return None


def _jwt_username(jwt: str) -> str | None:
    """Extract preferred_username from JWT payload without verification."""
    parts = jwt.split(".")
    if len(parts) < 2:
        return None
    padded = parts[1] + "=" * (-len(parts[1]) % 4)
    try:
        payload = json.loads(base64.urlsafe_b64decode(padded))
        name = payload.get("preferred_username")
        return name if isinstance(name, str) and name else None
    except Exception:
        return None


def _parse_ws_url(ws_url: str) -> tuple[str, str, str]:
    """Return (base_url, per_user_workstation, environment) from ws(s) URL.

    Expected format: ws(s)://host/workstation/<env>/<per_user_name>/ssh
    """
    scheme = "https" if ws_url.startswith("wss://") else "http"
    without_scheme = ws_url.split("://", 1)[1]
    host, _, path = without_scheme.partition("/")
    segments = [s for s in path.split("/") if s]
    # segments: ["workstation", env, per_user_name, "ssh"]
    if len(segments) >= 4 and segments[0] == "workstation" and segments[-1] == "ssh":
        return f"{scheme}://{host}", segments[2], segments[1]
    return f"{scheme}://{host}", "", ""


def _api_request(
    base_url: str,
    jwt: str,
    method: str,
    path: str,
    query: dict | None = None,
    body: dict | None = None,
    timeout: int = 30,
) -> dict:
    """Make an API call and return the parsed JSON response."""
    url = f"{base_url}{path}"
    if query:
        qs = "&".join(f"{k}={urllib.parse.quote(str(v))}" for k, v in query.items())
        url = f"{url}?{qs}"

    data = json.dumps(body).encode() if body else None
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("Cookie", f"{COOKIE_NAME}={jwt}")
    req.add_header("Accept", "application/json")
    if data:
        req.add_header("Content-Type", "application/json")

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as exc:
        try:
            return json.loads(exc.read())
        except Exception:
            return {"error": f"HTTP {exc.code}"}
    except Exception as exc:
        return {"error": str(exc)}


def _get_workstation_state(
    base_url: str, jwt: str, per_user_name: str, environment: str
) -> str | None:
    """Return the lowercase workstation state, or None if not found."""
    result = _api_request(
        base_url,
        jwt,
        "GET",
        "/api/job/get",
        query={"job_name": per_user_name, "environment": environment, "json": "true"},
    )
    if result.get("error"):
        return None
    data = result.get("result")
    if isinstance(data, dict):
        state = data.get("state")
        if isinstance(state, str):
            return state.lower()
        for key in ("jobs", "items", "data"):
            nested = data.get(key)
            if isinstance(nested, list) and nested:
                s = nested[0].get("state") if isinstance(nested[0], dict) else None
                if isinstance(s, str):
                    return s.lower()
    if isinstance(data, list) and data:
        s = data[0].get("state") if isinstance(data[0], dict) else None
        if isinstance(s, str):
            return s.lower()
    return None


def _start_workstation(
    base_url: str, jwt: str, template_name: str, environment: str
) -> str | None:
    """Start a workstation. Returns error string or None on success."""
    result = _api_request(
        base_url,
        jwt,
        "POST",
        "/api/job/start",
        body={"job_name": template_name, "environment": environment},
    )
    err = result.get("error", "")
    if err and "already running" not in str(err).lower():
        return str(err)
    return None


def _eprint(*args, **kwargs) -> None:
    print(*args, file=sys.stderr, **kwargs)


def _ensure_workstation_running(
    base_url: str,
    jwt: str,
    per_user_name: str,
    template_name: str,
    environment: str,
    timeout_s: int = 300,
    poll_s: int = 3,
) -> None:
    """Check workstation state; start it if needed and wait until running."""
    state = _get_workstation_state(base_url, jwt, per_user_name, environment)
    if state == "running":
        return

    _eprint(f"Workstation '{per_user_name}' is not running, starting it...")

    if state not in {"pending", "starting"}:
        err = _start_workstation(base_url, jwt, template_name, environment)
        if err:
            _eprint(f"Failed to start workstation: {err}")
            sys.exit(1)

    _eprint("In cloud environments provisioning may take up to several minutes.")

    use_tty = sys.stderr.isatty()
    phrases = random.sample(_LOADING_PHRASES, len(_LOADING_PHRASES))
    phrase_idx = 0
    dot_variants = ("", ".", "..", "...")
    dot_idx = 0
    phrase_interval = 6.0
    dot_interval = phrase_interval / len(dot_variants)
    next_poll = 0.0
    next_phrase = time.time() + phrase_interval
    next_dot = 0.0
    deadline = time.time() + timeout_s

    def _render() -> None:
        nonlocal dot_idx
        line = f"{phrases[phrase_idx]}{dot_variants[dot_idx % len(dot_variants)]}"
        dot_idx += 1
        if use_tty:
            _eprint(f"\r{line:<80}", end="", flush=True)
        else:
            _eprint(line)

    if use_tty:
        _render()
        next_dot = time.time() + dot_interval
    else:
        _eprint(phrases[phrase_idx])

    while time.time() < deadline:
        now = time.time()

        if now >= next_poll:
            state = _get_workstation_state(base_url, jwt, per_user_name, environment)
            if state == "running":
                if use_tty:
                    _eprint("\r" + " " * 80 + "\r", end="", flush=True)
                _eprint("Workstation is running.")
                return
            next_poll = now + poll_s

        if now >= next_phrase:
            phrase_idx = (phrase_idx + 1) % len(phrases)
            next_phrase = now + phrase_interval
            if use_tty:
                dot_idx = 0
                _render()
                next_dot = now + dot_interval
            else:
                _eprint(phrases[phrase_idx])

        if use_tty and now >= next_dot:
            _render()
            next_dot = now + dot_interval

        sleep_for = max(
            0.05, min(0.5, min(next_poll, next_phrase, deadline) - time.time())
        )
        time.sleep(sleep_for)

    if use_tty:
        _eprint()
    _eprint(f"Timed out after {timeout_s}s waiting for workstation to start.")
    sys.exit(1)


def main() -> None:
    if len(sys.argv) != 2:
        _eprint("Usage: ssh_proxy_command.py <ws_url>")
        sys.exit(2)

    cfg = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    jwt = cfg.get("jwt_cookie")
    if not jwt:
        _eprint("No jwt_cookie found. Run 'datatailr login'.")
        sys.exit(1)

    ws_url = sys.argv[1]

    base_url, per_user_name, environment = _parse_ws_url(ws_url)
    if per_user_name and environment:
        username = _jwt_username(jwt)
        if username:
            suffix = f"-{username}"
            if per_user_name.endswith(suffix):
                template_name = per_user_name[: -len(suffix)]
                _ensure_workstation_running(
                    base_url, jwt, per_user_name, template_name, environment
                )

    websocat_bin = _resolve_websocat()
    if not websocat_bin:
        _eprint("websocat not found. Run 'datatailr setup-ssh' again.")
        sys.exit(1)

    os.execv(
        websocat_bin,
        [websocat_bin, "--binary", f"-H=Cookie: {COOKIE_NAME}={jwt}", ws_url],
    )


if __name__ == "__main__":
    main()
