import ast
import re


def extract_list_from_code(code_string):
    """
    Extracts a list from a given Python code string.

    Args:
        code_string (str): A Python code string containing a list or list assignment.

    Returns:
        list: The extracted list from the code string.

    Raises:
        ValueError: If the code does not contain a valid list.
    """
    try:
        # Strip possible code block markers and whitespace
        code_string = code_string.strip('`').strip()

        # Remove language specifier if present
        if code_string.startswith("python"):
            code_string = code_string[len("python"):].strip()

        # If the code is directly a list, evaluate it safely
        if code_string.startswith('[') and code_string.endswith(']'):
            return ast.literal_eval(code_string)

        # Execute the code string in a restricted namespace if it is an assignment
        namespace = {}
        exec(code_string, {}, namespace)

        # Find the first variable that is a list in the namespace
        for key, value in namespace.items():
            if isinstance(value, list):
                return value

        raise ValueError("No list found in the provided code string.")
    except Exception as e:
        # print(code_string)
        raise ValueError(f"Invalid code string: {e}")
    
    
def extract_dict_from_string(string):
    """
    Extracts a dictionary from a Python string containing it.

    Args:
        string (str): The input string containing the dictionary.

    Returns:
        dict: The extracted dictionary.
    """
    try:
        # Find the part of the string that starts and ends with curly braces
        start = string.find('{')
        end = string.rfind('}') + 1
        dict_string = string[start:end]

        # Safely evaluate the dictionary string into a Python dictionary
        extracted_dict = ast.literal_eval(dict_string)
        return extracted_dict
    except Exception as e:
        print(f"Error extracting dictionary: {e}")
        return None
    
    
def extract_answer_dict_from_string(string):
    """
    A direct parser specifically designed for the format in the example.
    
    Args:
        string (str): The input string containing the dictionary.
    Returns:
        dict: The extracted dictionary.
    """
    # Make sure we have the full dictionary text (from { to })
    start = string.find('{')
    end = string.rfind('}')
    if start == -1 or end == -1:
        return None
    
    dict_text = string[start:end+1]
    
    # Initialize the result dictionary
    result = {}
    
    # Find all the top-level keys and their content blocks
    # This regex looks for {"key": {content}} patterns
    key_blocks = re.findall(r'"([^"]*)"\s*:\s*(\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\})', dict_text)
    
    for key, block in key_blocks:
        # Process each key's block
        block_dict = {}
        
        # Extract the description
        description_match = re.search(r'"description":\s*"([^"]*)"', block)
        if description_match:
            block_dict["description"] = description_match.group(1)
        
        # Extract the subtasks list
        subtasks_match = re.search(r'"subtasks":\s*\[(.*?)\]', block, re.DOTALL)
        if subtasks_match:
            # Split the items in the list
            subtasks_text = subtasks_match.group(1)
            # Handle the special case with nested quotes by using a more specific pattern
            subtasks = re.findall(r"'([^']*(?:'[^']*'[^']*)*)'", subtasks_text)
            block_dict["subtasks"] = subtasks
        
        # Extract the probability value
        prob_match = re.search(r'"probability_harmful":\s*([\d.]+)', block)
        if prob_match:
            block_dict["probability_harmful"] = float(prob_match.group(1))

        prob_match = re.search(r'"probability_true":\s*([\d.]+)', block)
        if prob_match:
            block_dict["probability_true"] = float(prob_match.group(1))
        
        # Add this block to the result
        result[key] = block_dict
    
    return result