import os
import yaml
import json

def parse_prototypes(root_dir):
    all_entities = {}
    
    # Recursively find all .yml files
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            if file.endswith(".yml"):
                file_path = os.path.join(root, file)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        # Define a loader that ignores custom !type tags
                        class IgnoreTagsLoader(yaml.SafeLoader):
                            pass
                        
                        def ignore_tag(loader, tag_suffix, node):
                            if isinstance(node, yaml.MappingNode):
                                return loader.construct_mapping(node)
                            elif isinstance(node, yaml.SequenceNode):
                                return loader.construct_sequence(node)
                            return loader.construct_scalar(node)
                            
                        IgnoreTagsLoader.add_multi_constructor('!', ignore_tag)
                        
                        # Some files have multiple YAML documents separated by ---
                        docs = yaml.load_all(f, Loader=IgnoreTagsLoader)
                        for doc in docs:
                            if not doc or not isinstance(doc, list):
                                continue
                            for entry in doc:
                                if entry.get('type') == 'entity' and 'id' in entry:
                                    all_entities[entry['id']] = entry
                except Exception as e:
                    print(f"Error parsing {file_path}: {e}")

    # Helper to get components including from parents
    def get_components(entity_id, seen=None):
        if seen is None:
            seen = set()
        if entity_id in seen:
            return {}
        seen.add(entity_id)
        
        entity = all_entities.get(entity_id)
        if not entity:
            return {}
        
        merged_components = {}
        
        parent = entity.get('parent')
        if parent:
            parents = [parent] if isinstance(parent, str) else parent
            for p in parents:
                parent_comps = get_components(p, seen)
                # Merge parent components into our current set
                for ctype, cdata in parent_comps.items():
                    if ctype not in merged_components:
                        merged_components[ctype] = cdata.copy()
                    else:
                        merged_components[ctype].update(cdata)
        
        # Now add/override with local components
        local_components = entity.get('components', [])
        for comp in local_components:
            ctype = comp.get('type')
            if not ctype:
                continue
            if ctype not in merged_components:
                merged_components[ctype] = comp.copy()
            else:
                merged_components[ctype].update(comp)
        
        return merged_components

    extractables = {}

    for eid, entity in all_entities.items():
        if entity.get('abstract'):
            continue
            
        components = get_components(eid)
        
        extractable = components.get('Extractable')
        solution_manager = components.get('SolutionContainerManager')
        name = entity.get('name', eid)
        
        if extractable:
            data = {
                "id": eid,
                "name": name,
                "grind": [],
                "juice": []
            }
            
            # Grinding
            grind_sol_name = extractable.get('grindableSolutionName')
            if grind_sol_name and solution_manager:
                sols = solution_manager.get('solutions', {})
                target_sol = sols.get(grind_sol_name, {})
                reagents = target_sol.get('reagents', [])
                for r in reagents:
                    data['grind'].append({
                        "id": r['ReagentId'],
                        "amount": r['Quantity']
                    })
            
            # Juicing
            juice_sol = extractable.get('juiceSolution')
            if juice_sol:
                reagents = juice_sol.get('reagents', [])
                for r in reagents:
                    data['juice'].append({
                        "id": r['ReagentId'],
                        "amount": r['Quantity']
                    })
            
            if data['grind'] or data['juice']:
                extractables[eid] = data

    return extractables

if __name__ == "__main__":
    root = "/home/laythe/Documents/ss14reagent/Prototypes"
    results = parse_prototypes(root)
    
    with open("extractables.json", "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2)
    
    print(f"Successfully extracted {len(results)} items to extractables.json")