import os
import ast

def discover_templates(root_dir="data/templates/branches"):
    """
    Recursively scans directory structure to find all template functions
    under each branch/domain/file.
    Returns:
        dict: {
            branch: {
                domain: {
                    file_path: [template_function_names]
                }
            }
        }
    """
    discovered = {}

    if not os.path.isdir(root_dir):
        print(f"Error: Directory '{root_dir}' not found.")
        return discovered

    # Top-level folders = branches
    for branch in sorted(os.listdir(root_dir)):
        branch_path = os.path.join(root_dir, branch)
        if os.path.isdir(branch_path):
            branch_dict = {}

            # Domains under branch
            for domain in sorted(os.listdir(branch_path)):
                domain_path = os.path.join(branch_path, domain)
                if os.path.isdir(domain_path):
                    domain_templates = {}

                    # Walk all Python files under domain
                    for dirpath, _, filenames in os.walk(domain_path):
                        for filename in sorted(filenames):
                            if filename.endswith(".py"):
                                file_path = os.path.join(dirpath, filename)
                                try:
                                    with open(file_path, "r", encoding="utf-8") as f:
                                        tree = ast.parse(f.read())
                                    template_functions = [
                                        node.name for node in ast.walk(tree)
                                        if isinstance(node, ast.FunctionDef) and node.name.startswith("template_")
                                    ]
                                    if template_functions:
                                        rel_path = os.path.relpath(file_path, root_dir)
                                        domain_templates[rel_path] = sorted(template_functions)
                                except Exception as e:
                                    print(f"Error parsing {file_path}: {e}")

                    if domain_templates:
                        branch_dict[domain] = domain_templates

            if branch_dict:
                discovered[branch] = branch_dict

    return discovered