# ssh_manager/cli.py

import os
import asyncio
from typing import Callable, Dict, Optional
from functools import wraps
from .utils import print_info, print_error, print_warning, Colors, safe_input
from .config import SSH_DIR, CONF_DIR, SOCKET_DIR, MAIN_CONFIG, DEFAULT_CONFIG_CONTENT
from .add_host import add_host
from .edit_host import edit_host
from .list_hosts import list_hosts
from .regen_key import regenerate_key
from .remove_host import remove_host

def ensure_ssh_setup() -> None:
    """
    Creates ~/.ssh, ~/.ssh/conf, and ~/.ssh/s if missing,
    and writes a default ~/.ssh/config if it doesn't exist.
    """
    directories = [
        (SSH_DIR, "SSH directory"),
        (CONF_DIR, "configuration directory"),
        (SOCKET_DIR, "socket directory")
    ]
    
    for directory, description in directories:
        if not os.path.isdir(directory):
            os.makedirs(directory, mode=0o700, exist_ok=True)
            print_info(f"Created {description}: {directory}")

    if not os.path.isfile(MAIN_CONFIG):
        with open(MAIN_CONFIG, "w") as f:
            f.write(DEFAULT_CONFIG_CONTENT)
        print_info(f"Created default SSH config at: {MAIN_CONFIG}")

def async_handler(func: Callable) -> Callable:
    """Decorator to handle async functions in the command dispatch"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        return asyncio.run(func(*args, **kwargs))
    return wrapper

class SSHManager:
    def __init__(self):
        self.commands: Dict[str, tuple[Callable, str]] = {
            "1": (self.list_hosts, "List Hosts"),
            "2": (self.add_host, "Add a Host"),
            "3": (self.edit_host, "Edit a Host"),
            "4": (self.regenerate_key, "Regenerate Key"),
            "5": (self.remove_host, "Remove Host"),
            "6": (self.exit_app, "Exit")
        }

    @async_handler
    async def list_hosts(self) -> None:
        await list_hosts(CONF_DIR)

    def add_host(self) -> None:
        add_host(CONF_DIR)

    @async_handler
    async def edit_host(self) -> None:
        await edit_host(CONF_DIR)

    @async_handler
    async def regenerate_key(self) -> None:
        await regenerate_key(CONF_DIR)

    @async_handler
    async def remove_host(self) -> None:
        await remove_host(CONF_DIR)

    def exit_app(self) -> None:
        print_info("Exiting...")
        raise SystemExit(0)

    def display_menu(self) -> None:
        print("\n" + f"{Colors.CYAN}{Colors.BOLD}SSH Config Manager Menu{Colors.RESET}")
        for key, (_, description) in self.commands.items():
            print(f"{key}. {description}")

    def handle_command(self, choice: str) -> bool:
        if choice not in self.commands:
            print_error("Invalid choice. Please select 1 through 6.")
            return True
        
        try:
            self.commands[choice][0]()
            return choice != "6"
        except SystemExit:
            return False
        except Exception as e:
            print_error(f"Error executing command: {str(e)}")
            return True

def main() -> int:
    ensure_ssh_setup()
    manager = SSHManager()
    
    # Display the server list on first load
    manager.list_hosts()

    while True:
        manager.display_menu()
        choice = safe_input("Select an option (1-6): ")
        if choice is None:
            continue
        
        if not manager.handle_command(choice.strip()):
            break

    return 0