import json
import secrets
import string
from collections import defaultdict

def generate_secure_token(length=20):
    charset = string.ascii_letters + string.digits + "!@#$%^&*()[]{}=+-_~"
    return ''.join(secrets.choice(charset) for _ in range(length))

def process_scratch_blocks(all_generated_blocks, generated_output_json):
    processed_blocks = {}

    # Initialize dictionaries to store and reuse generated unique IDs
    # This prevents creating multiple unique IDs for the same variable/broadcast across different blocks
    variable_id_map = defaultdict(lambda: generate_secure_token(20))
    broadcast_id_map = defaultdict(lambda: generate_secure_token(20))

    for block_id, gen_block_data in generated_output_json.items():
        processed_block = {}
        all_gen_block_data = all_generated_blocks.get(block_id, {})

        # Copy and update fields, inputs, next, parent, shadow, topLevel, mutation, and opcode
        processed_block["opcode"] = all_gen_block_data.get("op_code", gen_block_data.get("op_code"))
        processed_block["inputs"] = {}
        processed_block["fields"] = {}
        processed_block["shadow"] = all_gen_block_data.get("shadow", gen_block_data.get("shadow"))
        processed_block["topLevel"] = all_gen_block_data.get("topLevel", gen_block_data.get("topLevel"))
        processed_block["parent"] = all_gen_block_data.get("parent", gen_block_data.get("parent"))
        processed_block["next"] = all_gen_block_data.get("next", gen_block_data.get("next"))
        if "mutation" in all_gen_block_data:
            processed_block["mutation"] = all_gen_block_data["mutation"]

        # Process inputs
        if "inputs" in all_gen_block_data:
            for input_name, input_data in all_gen_block_data["inputs"].items():
                if input_name in ["SUBSTACK", "CONDITION"]:
                    # These should always be type 2
                    if isinstance(input_data, list) and len(input_data) == 2:
                        processed_block["inputs"][input_name] = [2, input_data[1]]
                    elif isinstance(input_data, dict) and input_data.get("kind") == "block":
                        processed_block["inputs"][input_name] = [2, input_data.get("block")]
                    else: # Fallback for unexpected formats, try to use the original if possible
                        processed_block["inputs"][input_name] = gen_block_data["inputs"].get(input_name, [2, None])

                elif isinstance(input_data, dict):
                    if input_data.get("kind") == "value":
                        # Case 1: Direct value input
                        processed_block["inputs"][input_name] = [
                            1,
                            [
                                4,
                                str(input_data.get("value", ""))
                            ]
                        ]
                    elif input_data.get("kind") == "block":
                        # Case 3: Nested block input
                        existing_shadow_value = ""
                        if input_name in gen_block_data.get("inputs", {}) and \
                           isinstance(gen_block_data["inputs"][input_name], list) and \
                           len(gen_block_data["inputs"][input_name]) > 2 and \
                           isinstance(gen_block_data["inputs"][input_name][2], list) and \
                           len(gen_block_data["inputs"][input_name][2]) > 1:
                            existing_shadow_value = gen_block_data["inputs"][input_name][2][1]
                        
                        processed_block["inputs"][input_name] = [
                            3,
                            input_data.get("block", ""),
                            [
                                10,  # Assuming 10 for number/string shadow
                                existing_shadow_value
                            ]
                        ]
                    elif input_data.get("kind") == "menu":
                        # Handle menu inputs like in event_broadcast
                        menu_option = input_data.get("option", "")
                        
                        # Generate or retrieve a unique ID for the broadcast message
                        broadcast_id = broadcast_id_map[menu_option] # Use defaultdict for unique IDs

                        processed_block["inputs"][input_name] = [
                            1,
                            [
                                11, # This is typically the code for menu dropdowns
                                menu_option,
                                broadcast_id
                            ]
                        ]
                elif isinstance(input_data, list):
                    # For cases like TOUCHINGOBJECTMENU, where input_data is a list [1, "block_id"]
                    processed_block["inputs"][input_name] = input_data


        # Process fields
        if "fields" in all_gen_block_data:
            for field_name, field_value in all_gen_block_data["fields"].items():
                if field_name == "VARIABLE" and isinstance(field_value, list) and len(field_value) > 0:
                    # Generate or retrieve a unique ID for the variable
                    variable_name = field_value[0]
                    unique_id = variable_id_map[variable_name] # Use defaultdict for unique IDs

                    processed_block["fields"][field_name] = [
                        variable_name,
                        unique_id
                    ]
                elif field_name == "STOP_OPTION":
                    processed_block["fields"][field_name] = [
                        field_value[0],
                        None
                    ]
                elif field_name == "TOUCHINGOBJECTMENU":
                    referenced_menu_block_id = all_gen_block_data["inputs"].get("TOUCHINGOBJECTMENU", [None, None])[1]
                    if referenced_menu_block_id and referenced_menu_block_id in all_generated_blocks:
                        menu_block = all_generated_blocks[referenced_menu_block_id]
                        menu_value = menu_block.get("fields", {}).get("TOUCHINGOBJECTMENU", ["", None])[0]
                        processed_block["fields"][field_name] = [menu_value, None]
                    else:
                        processed_block["fields"][field_name] = [field_value[0], None]
                else:
                    processed_block["fields"][field_name] = field_value
        
        # Remove unwanted keys from the processed block
        keys_to_remove = ["functionality", "block_shape", "id", "block_name", "block_type"]
        for key in keys_to_remove:
            if key in processed_block:
                del processed_block[key]
        
        processed_blocks[block_id] = processed_block
    return processed_blocks
    
# Path to your JSON file
if __name__ == "__main__":
    all_generated_blocks_path = 'all_generated_blocks.json'
    generated_output_json_path = 'generated_output_json.json'

    # Open and load the JSON files into Python dictionaries
    with open(all_generated_blocks_path, 'r') as f:
        all_generated_blocks_data = json.load(f)
        
    with open(generated_output_json_path, 'r') as f:
        generated_output_json_data = json.load(f)
        
    processed_blocks = process_scratch_blocks(all_generated_blocks_data, generated_output_json_data)
    print(json.dumps(processed_blocks, indent=4))