import xml.etree.ElementTree as ET
import sys
import os
from colorama import Fore, Style, init

# Initialisation de colorama pour Windows (ou autres plateformes)
init(autoreset=True)

def create_tree(children: list, dir):
    for child in children:
        if child.get('type') == 'Folder':
            path = os.path.join(dir, child.get('name'))
            if not os.path.exists(path):
                os.makedirs(path, exist_ok=True)
                print(f"{Fore.BLUE}[FOLDER CREATED] {Style.RESET_ALL} {path}")

            create_tree(child.get('children'), path)
        
        if child.get('type') == 'Script':
            name = child.get('name')
            if child.get('children'):
                new_dir = os.path.join(dir, name)
                if not os.path.exists(new_dir):
                    os.makedirs(new_dir, exist_ok=True)
                    print(f"{Fore.BLUE}[SCRIPT FOLDER CREATED] {Style.RESET_ALL} {new_dir}")

                name = 'init'
                
                with open(os.path.join(new_dir, f"{name}.lua"), 'w+') as f:
                    f.write(child.get('script'))
                    print(f"{Fore.GREEN}[SCRIPT FILE CREATED] {Style.RESET_ALL} {os.path.join(new_dir, f'{name}.lua')}")
                
                create_tree(child.get('children'), new_dir)
            else:
                if name == child.get('parent'):
                    name = 'init'
                
                with open(os.path.join(dir, f"{name}.lua"), 'w+') as f:
                    f.write(child.get('script'))
                    print(f"{Fore.GREEN}[SCRIPT FILE CREATED] {Style.RESET_ALL} {os.path.join(dir, f'{name}.lua')}")

                create_tree(child.get('children'), dir)


def get_children(parent: ET.Element):
    name = ''
    children = []

    for child in parent:
        properties = child.find('Properties')
        if child.tag == "Properties":
            name = child.findall('string')[0].text

        if child.tag == "Item":
            children.append({
                'type': child.attrib['class'],
                'name': properties.findall('string')[0].text,
                'script': child.attrib['class'] == 'Script' and properties.find('ProtectedString').text or None,
                'parent': name,
                'children': get_children(child),
            })

    return children

def main():
    if len(sys.argv) < 2:
        print(f"{Fore.RED}[ERROR] No file provided. Usage: rbxmx-extractor <filename>.rbxmx{Style.RESET_ALL}")
    else:
        name = sys.argv[1]

        print(f"{Fore.CYAN}[INFO] Starting parsing of the file: {name}{Style.RESET_ALL}")
        try:
            tree = ET.parse(name)
        except FileNotFoundError:
            print(f"{Fore.RED}[ERROR] File not found: {name}{Style.RESET_ALL}")
            return
        except ET.ParseError as e:
            print(f"{Fore.RED}[ERROR] Failed to parse XML: {e}{Style.RESET_ALL}")
            return

        root = tree.getroot()
        children = get_children(root)
        
        output_dir = 'out'
        os.makedirs(output_dir, exist_ok=True)
        print(f"{Fore.YELLOW}[INFO] Output directory created: {output_dir}{Style.RESET_ALL}")
        
        create_tree(children, output_dir)
        print(f"{Fore.CYAN}[INFO] Parsing and extraction completed! Check the 'out' folder.{Style.RESET_ALL}")