import re
import json
import pandas as pd

def expand_udt_references(db_struct, udts):
    """
    Recursively expand UDT references in the given DB structure using the UDT definitions.
    This function specifically expands fields designated as 'type' which reference UDTs.
    """
    if isinstance(db_struct, dict):
        for key, value in list(db_struct.items()):
            if isinstance(value, dict):
                # Recurse into dictionaries
                expand_udt_references(value, udts)
            elif isinstance(value, str) and key == "type":  # Only expand 'type' fields
                type_name = value.strip(
                    '"'
                )  # Remove quotes which may wrap UDT names with spaces
                if type_name in udts:
                    # Replace the UDT reference with its definition, if it exists
                    db_struct["fields"] = udts[
                        type_name
                    ].copy()  # Assume structure to insert is under 'fields'
                    print(f"Expanded UDT '{type_name}' at field '{key}' ")
    elif isinstance(db_struct, list):
        for item in db_struct:
            expand_udt_references(item, udts)


def handle_array_types(db_struct):
    """
    Handle array types once all UDTs are expanded.
    This function modifies the structure in place.
    """
    if isinstance(db_struct, dict):
        for key, value in list(db_struct.items()):
            if isinstance(value, dict):
                handle_array_types(value)
            elif isinstance(value, str):
                # Parsing array definitions, e.g., "Array[1..3] of Real"
                match = re.match(r"Array\[(\d+)\.\.(\d+)\] of (\w+)", value)
                if match:
                    lower_bound, upper_bound, base_type = (
                        int(match.group(1)),
                        int(match.group(2)),
                        match.group(3),
                    )
                    # Expand this field into multiple fields
                    db_struct.pop(key)  # Remove the original field
                    for i in range(lower_bound, upper_bound + 1):
                        db_struct[f"{key}_{i}"] = {"type": base_type, 'is_array_element': True}
                    print(
                        f"Expanded field '{key}' into array fields: {key}_{lower_bound} to {key}_{upper_bound} of type {base_type}"
                    )


type_sizes = {
    "Int": 2,
    "DInt": 4,
    "Word": 2,
    "Real": 4,
    "Bool": 2,  # We'll adjust this dynamically based on context (1 byte if alone, 1 bit if grouped)
    "String": 1,  # This will be multiplied by the specified size in brackets [n]
}


def calculate_offsets(db_struct, current_offset=0, parent=None):
    """
    Recursively calculate byte offsets for each field in the DB structure considering special types.
    """
    last_key_was_bool = False
    if isinstance(db_struct, dict):
        for key, value in list(db_struct.items()):
            if isinstance(value, dict):
                if "type" in value:
                    type_name = value["type"]
                    is_array_element = value.get('is_array_element', False)
                    size = type_sizes.get(
                        type_name, 0
                    )  # Default to 1 byte if type is not recognized

                    if not is_array_element and current_offset % 2 != 0:
                        current_offset += 1  # Align to the next even offset if it's not an array element
                        
                    # Special handling for String types
                    if "String" in type_name:
                        match = re.match(r"String\[(\d+)\]", type_name)
                        if match:
                                    length = int(match.group(1))
                                    size = length + 2  # Account for null-termination and string length prefix
                        else:
                            size = type_sizes.get(type_name, 1)  # Default to generic size if not an array

                    # Adjusting Bool sizes based on grouping
                    if type_name == "Bool":
                        if last_key_was_bool:  # This is a grouped bool
                            size = 0.125  # One bit per Bool if grouped
                        else:
                            size = 2  # Bools use a full byte if not grouped
                        last_key_was_bool = True
                    else:
                        last_key_was_bool = False

                    value["offset"] = current_offset
                    current_offset += size

                current_offset = calculate_offsets(
                    value, current_offset, value
                )  # Recurse into nested structs
    elif isinstance(db_struct, list):
        for item in db_struct:
            current_offset = calculate_offsets(item, current_offset, parent)
    return current_offset


def expand_dbs(udts, dbs):
    """
    Expand all UDT references in all DBs and then handle array types.
    """
    for db_name, db_content in dbs.items():
        print(f"Expanding DB: {db_name}")
        expand_udt_references(db_content, udts)
        handle_array_types(db_content)
        calculate_offsets(db_content)
        print(f"Completed expansion for DB: {db_name}")