import toml
import os
import subprocess

CONFIG_FILE = '/etc/frp/frpc.toml'

def load_config():
    try:
        if not os.path.exists(CONFIG_FILE):
            return {}
        with open(CONFIG_FILE, 'r') as f:
            return toml.load(f)
    except toml.TomlDecodeError as e:
        print(f"Failed to parse TOML file: {e}")
        return {}
    except Exception as e:
        print(f"An error occurred: {e}")
        return {}

def save_config(config):
    try:
        with open(CONFIG_FILE, 'w') as f:
            toml.dump(config, f)
    except Exception as e:
        print(f"Failed to save configuration: {e}")

def list_proxies_menu(config, action="edit"):
    if 'proxies' not in config or not config['proxies']:
        print("No proxies available.")
        return None

    print("\nExisting proxies:")
    for i, proxy in enumerate(config['proxies']):
        print(f"{i + 1}. {proxy['name']}")

    choice = input(f"Select a proxy to {action} (number): ")

    try:
        index = int(choice) - 1
        if 0 <= index < len(config['proxies']):
            return config['proxies'][index]['name']
    except ValueError:
        print("Invalid selection. Please enter a valid number.")
    return None

def list_proxies(config):
    if 'proxies' not in config or not config['proxies']:
        print("No proxies available.")
        return

    proxies = config['proxies']
    print(f"{'Name':<15} {'Type':<10} {'Local IP':<15} {'Local Port':<10} {'Remote Port':<10}")
    print("-" * 60)
    for proxy in proxies:
        print(f"{proxy['name']:<15} {proxy['type']:<10} {proxy['localIP']:<15} {proxy['localPort']:<10} {proxy['remotePort']:<10}")

def get_non_empty_input(prompt):
    """Helper function to ensure non-empty input for required fields."""
    while True:
        try:
            input_str = input(prompt).strip()
            if input_str:
                return input_str
            print("This field is required.")
        except KeyboardInterrupt:
            print("\nOperation cancelled by user.")
            return None

def get_int_input(prompt):
    """Helper function to ensure valid integer input."""
    while True:
        try:
            return int(input(prompt).strip())
        except ValueError:
            print("Please enter a valid integer.")
        except KeyboardInterrupt:
            print("\nOperation cancelled by user.")
            return None

def add_proxy(config):
    name = get_non_empty_input("Enter proxy name: ")
    if name is None:
        return

    if any(proxy['name'] == name for proxy in config.get('proxies', [])):
        print(f"A proxy with the name {name} already exists.")
        return

    typ = None
    while typ is None:
        try:
            print("Select proxy type:")
            print("1. tcp")
            print("2. udp")
            typ_choice = input("Enter choice: ").strip()
            if typ_choice == '1':
                typ = 'tcp'
            elif typ_choice == '2':
                typ = 'udp'
            else:
                print("Invalid choice. Please enter 1 or 2.")
        except KeyboardInterrupt:
            print("\nOperation cancelled by user.")
            return

    local_ip = input("Enter local IP (default 127.0.0.1): ").strip() or "127.0.0.1"
    local_port = get_int_input("Enter local port: ")
    remote_port = get_int_input("Enter remote port: ")

    if local_port is None or remote_port is None:
        return

    new_proxy = {
        'name': name,
        'type': typ,
        'localIP': local_ip,
        'localPort': local_port,
        'remotePort': remote_port
    }

    if 'proxies' not in config:
        config['proxies'] = []

    config['proxies'].append(new_proxy)
    save_config(config)
    print(f"Added proxy: {new_proxy}")

def edit_proxy(config):
    name = list_proxies_menu(config, "edit")
    if not name:
        return

    for proxy in config['proxies']:
        if proxy['name'] == name:
            print(f"Editing proxy {name}:")

            curr_type = proxy['type']
            print(f"Current type: {curr_type}")
            new_typ = None
            while new_typ is None:
                try:
                    print("Select new proxy type (leave blank to keep current):")
                    print("1. tcp")
                    print("2. udp")
                    type_choice = input("Enter choice: ").strip()
                    if type_choice == '':
                        new_typ = curr_type
                    elif type_choice == '1':
                        new_typ = 'tcp'
                    elif type_choice == '2':
                        new_typ = 'udp'
                    else:
                        print("Invalid choice. Please select 1 or 2, or leave blank to keep current.")
                except KeyboardInterrupt:
                    print("\nOperation cancelled by user.")
                    return

            proxy['type'] = new_typ

            try:
                local_ip = input("Enter new local IP (leave blank to keep current): ").strip()
                if local_ip:
                    proxy['localIP'] = local_ip

                local_port = input("Enter new local port (leave blank to keep current): ").strip()
                if local_port:
                    proxy['localPort'] = int(local_port)

                remote_port = input("Enter new remote port (leave blank to keep current): ").strip()
                if remote_port:
                    proxy['remotePort'] = int(remote_port)
            except ValueError:
                print("Invalid input. Keeping current values.")
            except KeyboardInterrupt:
                print("\nOperation cancelled by user.")
                return

            save_config(config)
            print(f"Edited proxy: {proxy}")
            return

def remove_proxy(config):
    name = list_proxies_menu(config, "remove")
    if not name:
        return

    config['proxies'] = [proxy for proxy in config['proxies'] if proxy['name'] != name]
    save_config(config)
    print(f"Removed proxy with name {name}.")

def configure_server(config):
    server_addr = get_non_empty_input("Enter server address: ")
    if server_addr is None:
        return

    server_port = get_int_input("Enter server port: ")
    if server_port is None:
        return

    config['serverAddr'] = server_addr
    config['serverPort'] = server_port
    save_config(config)
    print("Server configuration updated.")

def show_or_configure_server(config):
    server_addr = config.get('serverAddr', 'Not configured')
    server_port = config.get('serverPort', 'Not configured')

    print("\nCurrent server configuration:")
    print(f"Address: {server_addr}")
    print(f"Port: {server_port}")

    choice = input("Do you want to update the server configuration? (y/n): ").strip().lower()
    if choice == 'y':
        configure_server(config)

def first_run_setup(config):
    if 'serverAddr' not in config or 'serverPort' not in config:
        print("First time setup:")
        configure_server(config)

def prompt_restart_service():
    """Prompt the user if they would like to restart the frpc systemd service."""
    while True:
        choice = input("Do you want to restart the frpc service? (y/n): ").strip().lower()
        if choice == 'y':
            try:
                subprocess.run(['systemctl', 'restart', 'frpc'], check=True)
                print("frpc service restarted successfully.")
            except subprocess.CalledProcessError as e:
                print(f"Failed to restart frpc service: {e}")
            break
        elif choice == 'n':
            print("frpc service restart skipped.")
            break
        else:
            print("Invalid choice. Please enter 'y' or 'n'.")

def main():
    try:
        config = load_config()

        first_run_setup(config)

        while True:
            print("\nSelect an operation:")
            print("1. List proxies")
            print("2. Add proxy")
            print("3. Edit proxy")
            print("4. Remove proxy")
            print("5. Show/Configure server address/port")
            print("6. Exit")
            choice = input("Enter choice: ").strip()

            if choice == '1':
                list_proxies(config)
            elif choice == '2':
                add_proxy(config)
            elif choice == '3':
                edit_proxy(config)
            elif choice == '4':
                remove_proxy(config)
            elif choice == '5':
                show_or_configure_server(config)
            elif choice == '6':
                prompt_restart_service()
                break
            else:
                print("Invalid choice. Please try again.")
    except KeyboardInterrupt:
        print("\nExited by user.")

if __name__ == '__main__':
    main()