#!/usr/bin/env python3

from kitty import draw_to_terminal, get_position, set_position, hide_cursor, show_cursor #, draw_animation

from PIL import Image

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

# Convert lat/lon to Cartesian coordinates
def latlon_to_xyz(lat, lon, 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):
    # Run traceroute command
    result = subprocess.run(['traceroute', '-n', target, '-q', '1', '-w', '1,3,10'], capture_output=True, text=True)
    hops = re.findall(r'\n\s*\d+\s+([\d.]+)', result.stdout)

    coords = []
    for ip in hops:
        try:
            response = 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

def main():

    parser = argparse.ArgumentParser(
                    prog='kglobe',
                    description='Traceroute on a globe',
                    epilog='Requires kitty graphics protocol support in terminal')

    parser.add_argument('target')
    parser.add_argument('-e', '--external', action='store_true')

    args = parser.parse_args()

    locations = traceroute(args.target)

    globe = pv.Sphere(radius=1.0, theta_resolution=120, phi_resolution=120,
                        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
    points_3d = [latlon_to_xyz(lat, lon) for lat, lon in locations]

    pl=pv.Plotter(off_screen=(not args.external))
    pl.add_mesh(globe, color='tan', smooth_shading=True, texture=tex, show_edges=False)

    for pt in points_3d:
        city_marker = pv.Sphere(center=pt, radius=0.02)
        pl.add_mesh(city_marker, color='blue')
    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)
        pl.add_mesh(line, color='red', line_width=2)

    hide_cursor()
    y, x = get_position()
    print('\n' * 25, end='')

    frames = []
    try:
        if not args.external:
            while True:
                pl.camera.Azimuth(1)
                image = pl.screenshot(transparent_background=True, window_size=(512, 512))
                frames.append(Image.fromarray(image))
                set_position(y-25, x)
                draw_to_terminal(Image.fromarray(image))
        else:
            pl.show()
    finally:
        show_cursor()

if __name__ == '__main__':
    main()