from pathlib import Path
import yaml
import hcl2
from collections import defaultdict
from rich.console import Console
from rich.tree import Tree
import typer
import re

# Define the root paths
INSTANCES_DIR = Path("config/instances")

def extract_node_name(hcl_path):
    text = hcl_path.read_text()
    match = re.search(r'node_name\s*=\s*"([^"]+)"', text)
    if match:
        return match.group(1)
    return None

# Function to extract cobbler_mgmt_classes and profiles from config.yaml
def extract_config_data(config_path):
    with config_path.open("r") as f:
        config = yaml.safe_load(f)
    return (
        config.get("cobbler_mgmt_classes", []),
        config.get("profiles", []),
    )

# Build a dictionary mapping node_name to instances and their metadata
def build_node_tree():
    tree_data = defaultdict(list)
    for instance_dir in INSTANCES_DIR.iterdir():
        if not instance_dir.is_dir() or instance_dir.name in {"template"}:
            continue

        config_path = instance_dir / "config.yaml"
        hcl_path = instance_dir / "terragrunt.hcl"

        if not config_path.exists() or not hcl_path.exists():
            continue

        node_name = extract_node_name(hcl_path)
        if not node_name:
            continue

        classes, profiles = extract_config_data(config_path)
        tree_data[node_name].append({
            "instance": instance_dir.name,
            "classes": classes,
            "profiles": profiles
        })
    return tree_data

# CLI using Typer
app = typer.Typer()
console = Console()

@app.command()
def show():
    data = build_node_tree()
    root = Tree("📦 [bold blue]Node Overview[/bold blue]")
    for node, instances in sorted(data.items()):
        node_branch = root.add(f"[bold green]{node}[/bold green]")
        for inst in sorted(instances, key=lambda x: x['instance']):
            inst_branch = node_branch.add(f"[cyan]{inst['instance']}[/cyan]")
            if inst['classes']:
                inst_branch.add(f"🛠️ classes: {', '.join(inst['classes'])}")
            if inst['profiles']:
                inst_branch.add(f"📋 profiles: {', '.join(inst['profiles'])}")
    console.print(root)

if __name__ == "__main__":
    app()