#!/usr/bin/env python3
"""Mullvad exit node manager for Tailscale on macOS.

Usage:
  mullvad on <city or country>   - connect to Mullvad exit node
  mullvad off                     - disconnect Mullvad exit node
  mullvad status                  - show current Mullvad status
"""

import subprocess
import sys
import re
import random


def _run_tailscale(args, timeout=30):
    """Run a tailscale command, suppressing version mismatch warnings."""
    result = subprocess.run(
        ["tailscale"] + args,
        capture_output=True,
        text=True,
        timeout=timeout,
    )
    output = result.stdout
    if result.returncode != 0:
        if result.stderr:
            print(result.stderr.strip(), file=sys.stderr)
        result.check_returncode()
    return output


def _run_tailscale_live(args, timeout=30):
    """Run tailscale command with live output (for up/down)."""
    result = subprocess.run(
        ["tailscale"] + args,
        timeout=timeout,
    )
    if result.returncode != 0:
        sys.exit(result.returncode)


def get_exit_nodes():
    """Get all available exit nodes including Mullvad ones."""
    output = _run_tailscale(["exit-node", "list"])
    nodes = []

    for line in output.split("\n"):
        line = line.strip()
        if not line:
            continue
        if line.startswith("#") or line.startswith("Warning"):
            continue
        if line.startswith("IP "):
            continue

        match = re.match(
            r"^(\S+)\s+(\S+\.ts\.net)\.?\s+(.+?)\s{2,}(.+?)\s{2,}(\S+)$",
            line,
        )
        if match:
            ip, hostname, country, city, status = match.groups()
            hostname = hostname.rstrip(".")
            nodes.append(
                {
                    "ip": ip,
                    "hostname": hostname,
                    "country": country.strip(),
                    "city": city.strip(),
                    "status": status,
                }
            )
    return nodes


def find_nodes(nodes, location):
    """Find Mullvad nodes matching location (case-insensitive on city/country/hostname)."""
    location_lower = location.lower()
    matches = []
    for node in nodes:
        is_mullvad = "mullvad" in node["hostname"].lower()
        if not is_mullvad:
            continue

        if (
            location_lower in node["city"].lower()
            or location_lower in node["country"].lower()
            or location_lower in node["hostname"].lower()
        ):
            matches.append(node)

    return matches


def get_current_exit_node():
    """Get the currently active exit node hostname, or None."""
    try:
        output = _run_tailscale(["exit-node", "list"], timeout=10)
    except subprocess.CalledProcessError:
        return None
    for line in output.split("\n"):
        if "selected" in line:
            parts = line.split()
            if len(parts) >= 2:
                return parts[1].rstrip(".")
            return line.strip()
    return None


def show_status():
    """Show current Mullvad status."""
    current = get_current_exit_node()
    if current:
        if "mullvad" in current.lower():
            print(f"Mullvad: ON  -> {current}")
        else:
            print(f"Exit node: {current} (not Mullvad)")
    else:
        print("Mullvad: OFF")


def cmd_on(location):
    """Connect to a Mullvad exit node matching the given location."""
    nodes = get_exit_nodes()
    matches = find_nodes(nodes, location)

    if not matches:
        print(f"No Mullvad exit nodes found for '{location}'")
        print("\nAvailable Mullvad locations:")
        shown = set()
        for node in nodes:
            is_mullvad = "mullvad" in node["hostname"].lower()
            if is_mullvad:
                key = f"{node['country']} - {node['city']}"
                if key not in shown:
                    print(f"  {key}")
                    shown.add(key)
        sys.exit(1)

    node = random.choice(matches)
    hostname = node["hostname"]

    print(f"Connecting to Mullvad: {node['country']} - {node['city']} ({hostname})")
    print("Tearing down Tailscale...")
    _run_tailscale_live(["down", "--accept-risk=lose-ssh"])
    print("Bringing up Tailscale with Mullvad exit node...")
    _run_tailscale_live(
        ["up", "--reset", "--accept-risk=lose-ssh", f"--exit-node={hostname}", "--exit-node-allow-lan-access=false", "--accept-routes"]
    )
    print(f"Mullvad: ON  -> {hostname}  ({node['country']}, {node['city']})")


def cmd_off():
    """Disconnect from the Mullvad exit node."""
    print("Tearing down Tailscale...")
    _run_tailscale_live(["down", "--accept-risk=lose-ssh"])
    print("Bringing up Tailscale without exit node...")
    _run_tailscale_live(["up", "--reset", "--accept-risk=lose-ssh", "--accept-routes"])
    print("Mullvad: OFF")


def main():
    if len(sys.argv) < 2:
        print("Usage: mullvad on <city or country> | mullvad off | mullvad status")
        sys.exit(1)

    cmd = sys.argv[1].lower()

    if cmd == "on":
        if len(sys.argv) < 3:
            print("Usage: mullvad on <city or country>")
            print("  e.g. mullvad on amsterdam")
            print("  e.g. mullvad on netherlands")
            sys.exit(1)
        cmd_on(sys.argv[2])

    elif cmd == "off":
        cmd_off()

    elif cmd == "status":
        show_status()

    else:
        print(f"Unknown command: {cmd}")
        print("Usage: mullvad on <city or country> | mullvad off | mullvad status")
        sys.exit(1)


if __name__ == "__main__":
    main()