import argparse
import ctypes
import subprocess
import sys


def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except Exception:
        return False


def decode_bytes(b):
    if not b:
        return ""
    # Use 'mbcs' to dynamically decode local windows terminal encoding, ignore bad bytes
    return b.decode('mbcs', errors='ignore')


def get_wsl_ip():
    try:
        # Run wsl hostname -I to get the ip
        # --exec bypasses the login shell minimizing noise
        result = subprocess.run(
            ["wsl.exe", "--exec", "hostname", "-I"],
            capture_output=True,
            check=True
        )
        output_str = decode_bytes(result.stdout)
        ips = output_str.strip().split()
        if ips:
            return ips[0]
        else:
            print("Error: Could not determine WSL IP address. Are you sure your WSL instance has an IP?")
            return None
    except subprocess.CalledProcessError as e:
        err_msg = decode_bytes(e.stderr) or decode_bytes(e.stdout)
        print(f"Error running wsl hostname -I: {err_msg}")
        return None
    except FileNotFoundError:
        print("Error: wsl.exe not found. Is WSL installed?")
        return None


def add_port(port):
    ip = get_wsl_ip()
    if not ip:
        return

    print(f"WSL IP detected: {ip}")
    print(f"Exposing port {port}...")

    # Add portproxy
    proxy_cmd = [
        "netsh", "interface", "portproxy", "add", "v4tov4",
        f"listenport={port}", "listenaddress=0.0.0.0",
        f"connectport={port}", f"connectaddress={ip}"
    ]
    try:
        subprocess.run(proxy_cmd, check=True, capture_output=True)
        print(f"Successfully added port proxy for port {port}")
    except subprocess.CalledProcessError as e:
        err_msg = decode_bytes(e.stderr) or decode_bytes(e.stdout)
        print(f"Failed to add port proxy: {err_msg}")
        return

    # Add firewall rule
    fw_cmd = [
        "netsh", "advfirewall", "firewall", "add", "rule",
        f"name=WSL Port {port}", "dir=in", "action=allow",
        "protocol=TCP", f"localport={port}"
    ]
    try:
        subprocess.run(fw_cmd, check=True, capture_output=True)
        print(f"Successfully added firewall rule 'WSL Port {port}'")
    except subprocess.CalledProcessError as e:
        err_msg = decode_bytes(e.stderr) or decode_bytes(e.stdout)
        print(f"Failed to add firewall rule: {err_msg}")


def remove_port(port):
    print(f"Removing exposed port {port}...")

    # Remove portproxy
    proxy_cmd = [
        "netsh", "interface", "portproxy", "delete", "v4tov4",
        f"listenport={port}", "listenaddress=0.0.0.0"
    ]
    try:
        subprocess.run(proxy_cmd, check=True, capture_output=True)
        print(f"Successfully removed port proxy for port {port}")
    except subprocess.CalledProcessError as e:
        err_msg = decode_bytes(e.stderr) or decode_bytes(e.stdout)
        print(f"Failed to remove port proxy: {err_msg}")

    # Remove firewall rule
    fw_cmd = [
        "netsh", "advfirewall", "firewall", "delete", "rule",
        f"name=WSL Port {port}"
    ]
    try:
        subprocess.run(fw_cmd, check=True, capture_output=True)
        print(f"Successfully removed firewall rule 'WSL Port {port}'")
    except subprocess.CalledProcessError as e:
        err_msg = decode_bytes(e.stderr) or decode_bytes(e.stdout)
        print(f"Failed to remove firewall rule: {err_msg}")


def list_ports():
    print("Currently active WSL port proxies:")
    try:
        # Show only v4tov4 portproxies
        subprocess.run(
            ["netsh", "interface", "portproxy", "show", "v4tov4"],
            check=True
        )
    except subprocess.CalledProcessError as e:
        err_msg = decode_bytes(e.stderr) or decode_bytes(e.stdout)
        print(f"Error checking port proxies: {err_msg}")


def main():
    parser = argparse.ArgumentParser(description="Tool to expose WSL ports to the LAN")
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Add command
    parser_add = subparsers.add_parser("add", help="Expose a port from WSL to the LAN")
    parser_add.add_argument("port", type=int, help="The port number to expose")

    # Remove command
    parser_remove = subparsers.add_parser("remove", help="Remove an exposed port")
    parser_remove.add_argument("port", type=int, help="The port number to remove")

    # List command
    subparsers.add_parser("list", help="List all currently active WSL port proxies")

    args = parser.parse_args()

    if args.command in ["add", "remove"] and not is_admin():
        print("ERROR: Administrative privileges are required to modify netsh port proxy and firewall rules.")
        print("Please run this terminal as Administrator and try again.")
        sys.exit(1)

    if args.command == "add":
        add_port(args.port)
    elif args.command == "remove":
        remove_port(args.port)
    elif args.command == "list":
        list_ports()
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
