import os
import sys
from collections import defaultdict

def find_test_scenarios(root_dir):
    """
    Recursively searches the given root directory for files ending with
    '.test_scenario.xml' and returns a dictionary mapping scenario names to their
    paths relative to the root directory.

    Args:
        root_dir (str): The absolute path to the starting directory (e.g., 'COMPONENTS').

    Returns:
        dict[str, str]: A dictionary mapping scenario names (without suffix) to 
                        their relative file paths.
    """
    if not os.path.isdir(root_dir):
        print(f"Error: Directory not found or not accessible: {root_dir}")
        return {} # Return empty dictionary

    print(f"Scanning directory: '{root_dir}'...")

    scenario_suffix = ".test_scenario.xml"
    
    # Dictionary comprehension: {scenario_name: relative_path}
    scenarios_map = {
        # Key: Scenario name (filename without suffix)
        filename.replace(scenario_suffix, ""): 
        # Value: Relative path
        os.path.relpath(os.path.join(dirpath, filename), root_dir)
        
        for dirpath, _, filenames in os.walk(root_dir)
        for filename in filenames if filename.endswith(scenario_suffix)
    }
    
    return scenarios_map

def organize_by_layer_component(scenarios_map):
    """
    Organizes scenario paths into a nested dictionary structure based on the file path:
    {Layer_Folder: {Component_Folder: [scenario_name, ...]}}
    
    It assumes the Layer is the first folder and the Component is the folder
    preceding the 'test' directory (i.e., the third-to-last segment).

    Args:
        scenarios_map (dict[str, str]): Dictionary mapping scenario names to their 
                                        relative file paths.

    Returns:
        defaultdict: Nested dictionary (Layer -> Component -> List of Scenario Names).
    """
    organized_data = defaultdict(lambda: defaultdict(list))
    
    # Iterate over the scenario name and path
    for scenario_name, path in scenarios_map.items():
        # Split path into segments using the OS separator
        segments = path.split(os.sep)
        
        # Layer is the first segment (e.g., 'application_layer', 'drivers')
        layer = segments[0]
        
        # Component is the third-to-last segment (e.g., 'actuator_manager', 'ammonia')
        # We assume the file is inside a 'test' folder inside a component folder.
        if len(segments) >= 3:
            component = segments[-3]
        else:
            # Fallback for scenarios found too close to the root
            component = "Root_Component" 

        # Populate the nested dictionary
        organized_data[layer][component].append(scenario_name)
        
    return organized_data

def scenario_scan(components_root_dir):
    """
    Main function to scan for test scenarios, print the organized structure, and
    return the resulting dictionaries.
    
    Returns:
        tuple[defaultdict, dict]: The organized layer/component structure and the 
                                  raw dictionary of scenario names to paths.
    """
    # 1. Find all relative paths (now a dictionary: {name: path})
    found_scenarios_map = find_test_scenarios(components_root_dir)

    if not found_scenarios_map:
        print(f"\nNo files ending with '.test_scenario.xml' were found in {components_root_dir}.")
        # Return empty structures if nothing is found
        return defaultdict(lambda: defaultdict(list)), {}
    
    num_scenarios = len(found_scenarios_map)
    
    # 2. Print the simple list of found paths
    print(f"\n--- Found {num_scenarios} Test Scenarios ---")
    for scenario_name, path in found_scenarios_map.items():
        print(f"Scenario: '{scenario_name}' | Relative Path: {os.path.join("components",path)}")

    # 3. Organize into the layer/component structure
    organized_scenarios = organize_by_layer_component(found_scenarios_map)

    # 4. Print the organized structure
    print("\n--- Organized Layer/Component Structure ---")
    for layer, components in organized_scenarios.items():
        print(f"\n[LAYER] {layer.upper()}:")
        for component, scenarios in components.items():
            scenario_list = ", ".join(scenarios)
            print(f"  [Component] {component}: {scenario_list}")

    return organized_scenarios, found_scenarios_map

if __name__ == "__main__":
    # Check if the user provided the path as a command-line argument
    if len(sys.argv) < 2:
        print("Usage: python scan_scenarios.py <path_to_components_folder>")
        print("\nExample:")
        print("  python scan_scenarios.py \"D:\\ASF_01\\GITEA\\ASF-SH\\COMPONENTS\"")
        sys.exit(1)

    # The first argument is the directory path
    components_root_dir = sys.argv[1]
    
    # The return value from scenario_scan now includes the dictionary you requested
    organized_data, scenario_map = scenario_scan(components_root_dir)
    print(scenario_map)