#!/usr/bin/env python3
import json
import subprocess
import sys
import os
import getpass

def get_tunnels():
    try:
        result = subprocess.run(['ngrok', 'api', 'tunnels', 'list'], 
                              capture_output=True, text=True, check=True)
        tunnels = json.loads(result.stdout)
        return tunnels['tunnels']
    except Exception as e:
        print(f"Error getting tunnels: {e}")
        return []

def get_ssh_tunnel():
    tunnels = get_tunnels()
    for tunnel in tunnels:
        if tunnel['proto'] == 'tcp' and tunnel['forwards_to'] == 'localhost:22':
            url = tunnel['public_url']
            host_port = url.replace('tcp://', '')
            host, port = host_port.rsplit(':', 1)
            return host, port
    
    print("No SSH tunnel found")
    return None, None

def list_tunnels():
    tunnels = get_tunnels()
    if not tunnels:
        print("No tunnels found")
        return
    
    for tunnel in tunnels:
        url = tunnel['public_url']
        forwards_to = tunnel['forwards_to']
        
        if tunnel['proto'] == 'tcp':
            host_port = url.replace('tcp://', '')
            host, port = host_port.rsplit(':', 1)
            local_port = forwards_to.split(':')[-1]
            print(f"{host}:{port} -> localhost:{local_port}")
        else:
            domain = url.replace('https://', '').replace('http://', '')
            local_port = forwards_to.split(':')[-1]
            ngrok_port = "443" if url.startswith('https://') else "80"
            print(f"{domain}:{ngrok_port} -> localhost:{local_port}")

def main():
    if len(sys.argv) < 2:
        print("Usage: ncon [ssh|list]")
        sys.exit(1)
    
    if sys.argv[1] == 'ssh':
        host, port = get_ssh_tunnel()
        if host and port:
            cmd = ['ssh', '-p', port, f'david@{host}']
            print(f"Connecting to: {' '.join(cmd)}")
            os.execvp('ssh', cmd)
        else:
            print("Failed to find SSH tunnel")
            sys.exit(1)
    elif sys.argv[1] == 'list':
        list_tunnels()
    else:
        print("Unknown command. Use: ncon [ssh|list]")
        sys.exit(1)

if __name__ == '__main__':
    main()