import argparse
import requests
import json
import sys

BASE_URL = "http://localhost:8080"  # Update this if your server is running on a different host or port

def create_or_update_peer(public_key, allowed_ips):
    url = f"{BASE_URL}/peers"
    data = {
        "PublicKey": public_key,
        "AllowedIPs": allowed_ips
    }
    
    # First, try to update an existing peer
    response = requests.put(f"{url}/{public_key}", json=data)
    
    if response.status_code == 404:
        # If the peer doesn't exist, create a new one
        response = requests.post(url, json=data)
    
    if response.status_code in (200, 201):
        result = response.json()
        if "warning" in result:
            print(f"Warning: {result['warning']}")
        else:
            print(result['message'])
    else:
        print(f"Error: {response.status_code} - {response.text}")    

def delete_peer(public_key):
    url = f"{BASE_URL}/peers/{public_key}"
    response = requests.delete(url)
    
    if response.status_code == 200:
        print("Peer deleted successfully.")
    else:
        print(f"Error: {response.status_code} - {response.text}")

def list_peers():
    url = f"{BASE_URL}/peers"
    response = requests.get(url)
    
    if response.status_code == 200:
        peers = response.json()
        print(json.dumps(peers, indent=2))
    else:
        print(f"Error: {response.status_code} - {response.text}")

def restore_config():
    url = f"{BASE_URL}/restore"
    response = requests.post(url)
    
    if response.status_code == 200:
        print("Configuration restored successfully.")
    else:
        print(f"Error: {response.status_code} - {response.text}")

def main():
    parser = argparse.ArgumentParser(description="WireGuard Config Manager Client")
    parser.add_argument("action", choices=["create", "update", "delete", "list", "restore"], help="Action to perform")
    parser.add_argument("--public-key", help="Public key of the peer")
    parser.add_argument("--allowed-ips", help="Allowed IPs for the peer")
    
    args = parser.parse_args()
    
    if args.action in ["create", "update"]:
        if not args.public_key or not args.allowed_ips:
            print("Error: Both --public-key and --allowed-ips are required for create/update actions.")
            sys.exit(1)
        create_or_update_peer(args.public_key, args.allowed_ips)
    elif args.action == "delete":
        if not args.public_key:
            print("Error: --public-key is required for delete action.")
            sys.exit(1)
        delete_peer(args.public_key)
    elif args.action == "list":
        list_peers()
    elif args.action == "restore":
        restore_config()

if __name__ == "__main__":
    main()