import os
from collections import OrderedDict

class Colors:
    GREEN  = "\033[0;32m"
    RED    = "\033[0;31m"
    YELLOW = "\033[1;33m"
    CYAN   = "\033[0;36m"
    BOLD   = "\033[1m"
    RESET  = "\033[0m"

def print_error(message):
    print(f"{Colors.RED}{Colors.BOLD}[✖]  {Colors.RESET}{message}")

def print_warning(message):
    print(f"{Colors.YELLOW}{Colors.BOLD}[⚠]  {Colors.RESET}{message}")

def print_info(message):
    print(f"{Colors.GREEN}{Colors.BOLD}[✔]  {Colors.RESET}{message}")

def load_config_file(file_path):
    """
    Parse a single SSH config file and return a list of host blocks.
    Each block is an OrderedDict with keys like 'Host', 'HostName', etc.
    """
    blocks = []
    host_data = None

    try:
        with open(file_path, 'r') as f:
            lines = f.readlines()
    except Exception as e:
        print_error(f"Error reading SSH config file {file_path}: {e}")
        return blocks

    for line in lines:
        stripped_line = line.strip()
        # Skip empty lines and comments
        if not stripped_line or stripped_line.startswith('#'):
            continue

        # Start of a new Host block
        if stripped_line.lower().startswith('host '):
            host_labels = stripped_line.split()[1:]
            # Pick the first label that isn't a wildcard
            for label in host_labels:
                if '*' not in label:
                    # If we already have a host_data in progress, close it out
                    if host_data:
                        blocks.append(host_data)
                    host_data = OrderedDict({'Host': label})
                    break
        elif host_data:
            # Split on the first whitespace into key/value
            if ' ' in stripped_line:
                key, value = stripped_line.split(None, 1)
                host_data[key] = value.strip()

    # Add the last block if it exists
    if host_data:
        blocks.append(host_data)

    return blocks

def edit_host(CONF_DIR):
    """
    Let the user update fields for an existing host.
    1) Ask which host label to edit
    2) Locate its subdirectory + config
    3) Update (HostName, User, Port, IdentityFile) as needed
    4) Rewrite the config
    """
    host_label = input("Enter the Host label to edit: ").strip()
    if not host_label:
        print_error("Host label cannot be empty.")
        return

    host_dir = os.path.join(CONF_DIR, host_label)
    config_path = os.path.join(host_dir, "config")
    if not os.path.isfile(config_path):
        print_warning(f"No config file found at {config_path}; cannot edit this host.")
        return

    # Load the config file and look for the relevant host block
    blocks = load_config_file(config_path)
    if not blocks:
        print_warning(f"No valid Host blocks found in {config_path}")
        return

    # We'll assume there's only one block in each config
    # (the "Host <label>"). If multiple blocks exist, adapt accordingly.
    target_block = None
    for b in blocks:
        if b.get("Host") == host_label:
            target_block = b
            break

    if not target_block:
        print_warning(f"No matching Host '{host_label}' found in {config_path}")
        return

    old_hostname = target_block.get("HostName", "")
    old_user     = target_block.get("User", "")
    old_port     = target_block.get("Port", "22")
    old_identity = target_block.get("IdentityFile", "")

    print_info("Leave a field blank to keep its current value.")
    new_hostname = input(f"Enter new HostName [{old_hostname}]: ").strip()
    new_user     = input(f"Enter new User [{old_user}]: ").strip()
    new_port     = input(f"Enter new Port [{old_port}]: ").strip()
    new_ident    = input(f"Enter new IdentityFile [{old_identity}]: ").strip()

    # If user leaves it blank, keep the old value
    final_hostname = new_hostname if new_hostname else old_hostname
    final_user     = new_user if new_user else old_user
    final_port     = new_port if new_port else old_port
    final_ident    = new_ident if new_ident else old_identity

    # Rebuild the config lines
    new_config_lines = [
        f"Host {host_label}",
        f"    HostName {final_hostname}",
        f"    User {final_user}",
        f"    Port {final_port}"
    ]
    if final_ident:
        new_config_lines.append(f"    IdentityFile {final_ident}")

    # Overwrite the file
    try:
        with open(config_path, "w") as f:
            for line in new_config_lines:
                f.write(line + "\n")
        print_info(f"Updated config at: {config_path}")
    except Exception as e:
        print_error(f"Failed to update config: {e}")