kglobe/kglobe.py

#!/usr/bin/env python3

import kitty

import pyvista as pv
import numpy as np
import math
import subprocess
import re
import requests
import argparse
import asyncio

from terminalplotter import TerminalPlotter

# TODO: Color arches based on latency
# TODO: Text info (num hops etc.)
# TODO: Mouse support
# DONE: Interactive globe (spin w/ keys)
# DONE: Image spacing
# DONE: Async rendering?

EXAMPLE_ROUTE = [
        (47.996, 7.849), # freiburg
        (50.110, 8.682), # ffm
        (52.231, 21.006), # warsaw
        (12.988, 77.622), # bangalore
        (22.350, 114.184), # hong kong
        (-33.869, 151.208) # sydney
]

# Convert lat/lon to Cartesian coordinates
def latlon_to_xyz(lat: float, lon: float, radius=1.0):
    lat_rad = np.radians(lat)
    lon_rad = np.radians(lon)
    x = radius * np.cos(lat_rad) * np.cos(lon_rad)
    y = radius * np.cos(lat_rad) * np.sin(lon_rad)
    z = radius * np.sin(lat_rad)
    return np.array([x, y, z])

# Create an arch between two 3D points
def generate_arch(p1, p2, height_factor=0.2, n_points=100):
    # Normalize input points to lie on the unit sphere
    p1 = p1 / np.linalg.norm(p1)
    p2 = p2 / np.linalg.norm(p2)

    # Compute angle between p1 and p2
    omega = np.arccos(np.clip(np.dot(p1, p2), -1, 1))
    if omega == 0:
        return np.tile(p1, (n_points, 1))  # degenerate case

    t = np.linspace(0, 1, n_points)
    sin_omega = np.sin(omega)

    # Spherical linear interpolation (slerp)
    arch_points = (
        np.sin((1 - t)[:, None] * omega) * p1[None, :]
        + np.sin(t[:, None] * omega) * p2[None, :]
    ) / sin_omega

    # Add radial height offset based on sine curve
    heights = 1 + np.sin(np.pi * t) * height_factor
    arch_points *= heights[:, None]  # Scale outward from center

    return arch_points


def traceroute(target: str) -> list[tuple[int, int]]:
    # Run traceroute command
    result = subprocess.run(
        ["traceroute", "-n", target, "-q", "1", "-w", "1,3,10"],
        capture_output=True,
        text=True,
    )
    hops: list[str] = re.findall(r"\n\s*\d+\s+([\d.]+)", result.stdout)

    coords: list[tuple[int, int]] = []
    for ip in hops:
        try:
            response: dict = requests.get(f"http://ip-api.com/json/{ip}").json()
            if response["status"] == "success":
                coords.append((response["lat"], response["lon"]))
        except Exception:
            continue
    return coords


async def main():

    parser = argparse.ArgumentParser(
        prog="kglobe",
        description="Traceroute on a globe",
        epilog="Requires kitty graphics protocol support in terminal",
    )
    parser.add_argument("-t", "--traceroute", default=None)
    parser.add_argument("--example", action="store_true")
    parser.add_argument("-s", "--size", nargs=2, type=int, help="width x height in px")
    args = parser.parse_args()

    locations = []
    if args.traceroute:
        locations = traceroute(args.traceroute)

    globe = pv.Sphere(
        radius=1.0,
        theta_resolution=60,
        phi_resolution=60,
        start_theta=270.001,
        end_theta=270,
    )

    tex = pv.examples.load_globe_texture()
    globe.active_texture_coordinates = np.zeros((globe.points.shape[0], 2))
    globe.active_texture_coordinates[:, 0] = 0.5 + np.arctan2(
        globe.points[:, 1], globe.points[:, 0]
    ) / (2 * math.pi)
    globe.active_texture_coordinates[:, 1] = (
        0.5 + np.arcsin(globe.points[:, 2]) / math.pi
    )

    # Convert to 3D coordinates
    if args.example:
        points_3d = [latlon_to_xyz(lat, lon) for lat, lon in EXAMPLE_ROUTE]
    else:
        points_3d = [latlon_to_xyz(lat, lon) for lat, lon in locations]

    if args.size:
        width, height = args.size
    else:
        height, width = await kitty.get_terminal_size_pixel()
    plotter = TerminalPlotter(width, height)
    plotter.add_mesh(globe, color="tan", smooth_shading=False, texture=tex, show_edges=False)

    for point in points_3d:
        city_marker = pv.Sphere(center=point, radius=0.02)
        plotter.add_mesh(city_marker, color="blue")

    labels = [str(_) for _ in range(len(points_3d))]
    raised_points = [ point * 1.1 for point in points_3d]
    if raised_points:
        plotter.add_point_labels(raised_points, labels, point_size=0, font_size=14)

    for i in range(len(points_3d[:-1])):
        arch = generate_arch(points_3d[i], points_3d[i + 1], height_factor=0.2)
        line = pv.lines_from_points(arch, close=False)
        plotter.add_mesh(line, color="red", line_width=2)


    kitty.hide_cursor()
    try:
        #asyncio.run(plotter.run())
        await plotter.run()
    finally:
        kitty.show_cursor()


if __name__ == "__main__":
    asyncio.run(main())