import json
import argparse
import sys
import os

def escape_label(label):
    """Escapes strings for Mermaid labels."""
    if isinstance(label, str):
        return label.replace('"', '&quot;').replace('(', '&#40;').replace(')', '&#41;')
    return str(label)

def generate_mermaid(data):
    """Generates Mermaid code from the state dictionary."""
    lines = ["graph TD"]
    
    # Styles
    lines.append("    %% Styles")
    lines.append("    classDef register fill:#e1bee7,stroke:#4a148c,stroke-width:2px;")
    lines.append("    classDef valueNode fill:#b2dfdb,stroke:#004d40,stroke-width:2px;")
    lines.append("    classDef rootNode fill:#ffccbc,stroke:#bf360c,stroke-width:4px;")

    nodes = data.get("nodes", {})
    registers = data.get("registers", {})

    # Generate Nodes
    for node_id, node_data in nodes.items():
        value = node_data.get("value", "")
        escaped_value = escape_label(value)
        
        # Determine class
        node_class = "valueNode"
        if node_id == "root":
            node_class = "rootNode"
            
        # Node definition: id["<b>id</b><br/>value"]
        lines.append(f'    {node_id}["<b>{node_id}</b><br/>{escaped_value}"]:::{node_class}')
        
        # Edges from this node
        edges = node_data.get("edges", {})
        for edge_label, target_id in edges.items():
            lines.append(f'    {node_id} -->|{edge_label}| {target_id}')

    # Generate Registers
    lines.append("    subgraph Registers")
    for reg_id, reg_value in registers.items():
        # Check if register value is a node ID
        is_node_ref = isinstance(reg_value, str) and reg_value in nodes
        
        reg_node_id = f"reg_{reg_id}"
        
        if is_node_ref:
            # Register points to a node
            lines.append(f'        {reg_node_id}("Reg {reg_id}"):::register')
        else:
            # Register has an immediate value
            lines.append(f'        {reg_node_id}("Reg {reg_id}: {reg_value}"):::register')
            
    lines.append("    end")
    
    # Edges from registers to nodes (outside subgraph to avoid strict nesting issues if we want)
    # Actually Mermaid handles edges cross-subgraph fine.
    for reg_id, reg_value in registers.items():
        if isinstance(reg_value, str) and reg_value in nodes:
             reg_node_id = f"reg_{reg_id}"
             lines.append(f'    {reg_node_id} -.-> {reg_value}')

    return "\n".join(lines)

def main():
    parser = argparse.ArgumentParser(description="Convert gASM state JSON to Mermaid graph.")
    parser.add_argument("input_file", help="Path to the input JSON file")
    parser.add_argument("output_file", nargs="?", help="Path to the output Mermaid file (default: stdout)")
    
    args = parser.parse_args()

    try:
        with open(args.input_file, 'r') as f:
            data = json.load(f)
            
        mermaid_output = generate_mermaid(data)
        
        if args.output_file:
            with open(args.output_file, 'w') as f:
                f.write(mermaid_output)
            print(f"Successfully wrote Mermaid graph to {args.output_file}")
        else:
            print(mermaid_output)
            
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()
