# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def validate_node_definition(node_def: Dict) -> bool:
    """Validate node definition structure"""
    if not isinstance(node_def, dict):
        return False
    return all(
        isinstance(v, dict) and all(isinstance(k, str) for k in v.keys())
        for v in node_def.values()
    )

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def get_node_definitions(chain, structure: str, example: Dict) -> Dict[str, Dict[str, str]]:
    """Get node definitions with retry logic"""
    try:
        # Get response from LLM
        response = chain.invoke({"structure": structure, "example": example})

        # Parse response
        node_defs = ast.literal_eval(response)

        # Validate structure
        if not validate_node_definition(node_defs):
            raise ValueError("Invalid node definition structure")

        return node_defs

    except (ValueError, SyntaxError) as e:
        logger.error(f"Error parsing node definitions: {e}")
        raise

# Updated node definition template
node_example = {
    "NodeLabel1": {"property1": "row['property1']", "property2": "row['property2']"},
    "NodeLabel2": {"property1": "row['property1']", "property2": "row['property2']"},
    "NodeLabel3": {"property1": "row['property1']", "property2": "row['property2']"},
}

define_nodes_prompt = PromptTemplate(
    input_variables=["example", "structure"],
    template=("""
        Analyze the dataset structure below and extract the entity labels for nodes and their properties.n
        The node properties should be based on the dataset columns and their values.n
        Return the result as a dictionary where the keys are the node labels and the values are the node properties.nn
        Example: {example}nn

        Dataset Structure:n{structure}nn

        Make sure to include all the possible node labels and their properties.n
        If a property can be its own node, include it as a separate node label.n
        Please do not report triple backticks to identify a code block, just return the list of tuples.n
        Return only the dictionary containing node labels and properties, and don't include any other text or quotation.

        """
    ),
)

# Execute with error handling
try:
    node_chain = define_nodes_prompt | llm

    node_definitions = get_node_definitions(node_chain, structure=node_structure, example=node_example)
    logger.info(f"Node Definitions: {node_definitions}")
except Exception as e:
    logger.error(f"Failed to get node definitions: {e}")
    raise