#!/usr/bin/env python3
from argparse import ArgumentParser
from os import environ
import shlex
import subprocess
import sys
from typing import List
from protonvpn_cli.cli_wrapper import CLIWrapper
from protonvpn_nm_lib.api import FeatureEnum
from protonvpn_nm_lib.core.servers.list import LogicalServer

cli = CLIWrapper()
parser = ArgumentParser()

args = sys.argv[1:]

if not args:
    args = shlex.split(environ.get("PVPN_CMD_ARGS") or "")

for alias in [
    ["-s", "--streaming"],
    ["--sc"],
    ["--p2p"],
    ["--tor"]
]:
    parser.add_argument(*alias, action="store_true")

group = parser.add_mutually_exclusive_group()

for alias in [
    ["-f", "--fastest"],
    ["-r", "--random"]
]:
    group.add_argument(*alias, action="store_true")

parser.add_argument("--cc")
parsedArgs = parser.parse_args()
features: List[FeatureEnum] = list()

if parsedArgs.streaming:
    features.append(FeatureEnum.STREAMING)
if parsedArgs.sc:
    features.append(FeatureEnum.SECURE_CORE)
if parsedArgs.p2p:
    features.append(FeatureEnum.P2P)
if parsedArgs.tor:
    features.append(FeatureEnum.TOR)

def match(server: LogicalServer):
    return (parsedArgs.cc is None or server.exit_country.lower() == parsedArgs.cc.lower()) and (
            all(feature in server.features for feature in features))

matches = cli.protonvpn.get_session().servers.filter(match)
server = matches.get_fastest_server() if parsedArgs.fastest else matches.get_random_server()
subprocess.run(["protonwire", "connect", server.data["Domain"]])