import yaml
import socket
import shutil
from rich.console import Console

# Inicializa o console do Rich para uma saída colorida
console = Console()

# --- CONFIGURAÇÃO ---
HOSTS_FILE = 'hosts.yml'
BACKUP_FILE = 'hosts.yml.bak'


def load_hosts_config(filename=HOSTS_FILE):
    """Carrega a configuração dos hosts a partir de um arquivo YAML."""
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            # Usamos RoundTripLoader para tentar preservar comentários e formatação
            return yaml.safe_load(file)
    except FileNotFoundError:
        console.print(f"[bold red]Erro: O arquivo de configuração '{filename}' não foi encontrado.[/bold red]")
        return None
    except yaml.YAMLError as e:
        console.print(f"[bold red]Erro de formatação no arquivo YAML '{filename}': {e}[/bold red]")
        return None

def save_hosts_config(data, filename=HOSTS_FILE):
    """Salva os dados de configuração de volta no arquivo YAML."""
    try:
        with open(filename, 'w', encoding='utf-8') as file:
            yaml.dump(data, file, default_flow_style=False, sort_keys=False)
        return True
    except Exception as e:
        console.print(f"[bold red]Erro ao salvar o arquivo '{filename}': {e}[/bold red]")
        return False

def main():
    """Função principal do script."""
    console.print(f"[bold cyan]Iniciando a resolução de hostnames no arquivo '{HOSTS_FILE}'...[/bold cyan]")
    
    hosts_data = load_hosts_config()
    if not hosts_data:
        return

    # 1. Fazer backup do arquivo original
    try:
        shutil.copy(HOSTS_FILE, BACKUP_FILE)
        console.print(f"[green]Backup do arquivo original criado em '{BACKUP_FILE}'[/green]")
    except Exception as e:
        console.print(f"[bold red]Não foi possível criar o arquivo de backup: {e}[/bold red]")
        return

    # 2. Iterar e resolver os hostnames
    updated_count = 0
    failed_count = 0

    for host_key, config in hosts_data.items():
        hostname = config.get('ip')
        display_name = config.get('displayName', host_key)

        # Pula se o campo 'ip' já for um IP válido ou estiver vazio
        if not hostname or hostname.replace('.', '').isdigit():
            continue

        console.print(f"\n[bold]Processando:[/bold] {display_name}")
        console.print(f"  - Resolvendo hostname: '{hostname}'...", end="")
        
        try:
            # Tenta obter o endereço IPv4
            ipv4_address = socket.gethostbyname(hostname)
            
            if ipv4_address != hostname:
                console.print(f" [bold green]OK![/bold green] -> {ipv4_address}")
                # Atualiza o dicionário com o novo IP
                config['ip'] = ipv4_address
                updated_count += 1
            else:
                console.print(f" [yellow]Já é um IP.[/yellow]")

        except socket.gaierror:
            console.print(f" [bold red]FALHOU.[/bold red] Não foi possível resolver o hostname.")
            failed_count += 1

    # 3. Salvar as alterações
    if updated_count > 0:
        if save_hosts_config(hosts_data):
            console.print(f"\n[bold green]Arquivo '{HOSTS_FILE}' atualizado com sucesso![/bold green]")
    else:
        console.print("\n[yellow]Nenhum hostname precisou ser atualizado.[/yellow]")

    console.print(f"\n[bold]Resumo:[/bold]")
    console.print(f"  - [green]{updated_count} hosts atualizados.[/green]")
    console.print(f"  - [red]{failed_count} falhas na resolução.[/red]")


if __name__ == '__main__':
    main()