entity_relationship_generator.py

# import graphviz
# import json
# from tempfile import NamedTemporaryFile
# import os

# def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
#     """
#     Generates an entity relationship diagram from JSON input.

#     Args:
#         json_input (str): A JSON string describing the entity relationship diagram structure.
#                           It must follow the Expected JSON Format Example below.

#     Expected JSON Format Example:
#     {
#       "entities": [
#         {
#           "name": "Customer",
#           "type": "strong",
#           "attributes": [
#             {"name": "customer_id", "type": "primary_key"},
#             {"name": "first_name", "type": "regular"},
#             {"name": "last_name", "type": "regular"},
#             {"name": "email", "type": "regular"},
#             {"name": "phone", "type": "multivalued", "multivalued": true},
#             {"name": "age", "type": "derived", "derived": true},
#             {"name": "address", "type": "composite", "composite": true}
#           ]
#         },
#         {
#           "name": "Order",
#           "type": "strong",
#           "attributes": [
#             {"name": "order_id", "type": "primary_key"},
#             {"name": "customer_id", "type": "foreign_key"},
#             {"name": "order_date", "type": "regular"},
#             {"name": "total_amount", "type": "regular"},
#             {"name": "status", "type": "regular"},
#             {"name": "shipping_address", "type": "composite", "composite": true}
#           ]
#         },
#         {
#           "name": "Product",
#           "type": "strong",
#           "attributes": [
#             {"name": "product_id", "type": "primary_key"},
#             {"name": "name", "type": "regular"},
#             {"name": "description", "type": "regular"},
#             {"name": "price", "type": "regular"},
#             {"name": "category_id", "type": "foreign_key"},
#             {"name": "stock_quantity", "type": "regular"}
#           ]
#         },
#         {
#           "name": "Category",
#           "type": "strong",
#           "attributes": [
#             {"name": "category_id", "type": "primary_key"},
#             {"name": "name", "type": "regular"},
#             {"name": "description", "type": "regular"},
#             {"name": "parent_category_id", "type": "foreign_key"}
#           ]
#         },
#         {
#           "name": "OrderItem",
#           "type": "weak",
#           "attributes": [
#             {"name": "order_id", "type": "primary_key"},
#             {"name": "product_id", "type": "primary_key"},
#             {"name": "quantity", "type": "regular"},
#             {"name": "unit_price", "type": "regular"},
#             {"name": "discount", "type": "regular"}
#           ]
#         },
#         {
#           "name": "Payment",
#           "type": "strong",
#           "attributes": [
#             {"name": "payment_id", "type": "primary_key"},
#             {"name": "order_id", "type": "foreign_key"},
#             {"name": "payment_method", "type": "regular"},
#             {"name": "amount", "type": "regular"},
#             {"name": "payment_date", "type": "regular"},
#             {"name": "transaction_id", "type": "regular"}
#           ]
#         },
#         {
#           "name": "Supplier",
#           "type": "strong",
#           "attributes": [
#             {"name": "supplier_id", "type": "primary_key"},
#             {"name": "company_name", "type": "regular"},
#             {"name": "contact_person", "type": "regular"},
#             {"name": "email", "type": "regular"},
#             {"name": "phone", "type": "multivalued", "multivalued": true},
#             {"name": "address", "type": "composite", "composite": true}
#           ]
#         },
#         {
#           "name": "Review",
#           "type": "weak",
#           "attributes": [
#             {"name": "customer_id", "type": "primary_key"},
#             {"name": "product_id", "type": "primary_key"},
#             {"name": "rating", "type": "regular"},
#             {"name": "comment", "type": "regular"},
#             {"name": "review_date", "type": "regular"}
#           ]
#         }
#       ],
#       "relationships": [
#         {
#           "name": "Places",
#           "type": "regular",
#           "entities": ["Customer", "Order"],
#           "cardinalities": {
#             "Customer": "1",
#             "Order": "N"
#           }
#         },
#         {
#           "name": "Contains",
#           "type": "identifying",
#           "entities": ["Order", "OrderItem"],
#           "cardinalities": {
#             "Order": "1",
#             "OrderItem": "N"
#           }
#         },
#         {
#           "name": "Includes",
#           "type": "regular",
#           "entities": ["OrderItem", "Product"],
#           "cardinalities": {
#             "OrderItem": "N",
#             "Product": "1"
#           }
#         },
#         {
#           "name": "BelongsTo",
#           "type": "regular",
#           "entities": ["Product", "Category"],
#           "cardinalities": {
#             "Product": "N",
#             "Category": "1"
#           }
#         },
#         {
#           "name": "ProcessedBy",
#           "type": "regular",
#           "entities": ["Order", "Payment"],
#           "cardinalities": {
#             "Order": "1",
#             "Payment": "1"
#           }
#         },
#         {
#           "name": "Supplies",
#           "type": "regular",
#           "entities": ["Supplier", "Product"],
#           "cardinalities": {
#             "Supplier": "N",
#             "Product": "M"
#           }
#         },
#         {
#           "name": "Writes",
#           "type": "weak",
#           "entities": ["Customer", "Review"],
#           "cardinalities": {
#             "Customer": "1",
#             "Review": "N"
#           }
#         },
#         {
#           "name": "ReviewsProduct",
#           "type": "weak",
#           "entities": ["Review", "Product"],
#           "cardinalities": {
#             "Review": "N",
#             "Product": "1"
#           }
#         },
#         {
#           "name": "Subcategory",
#           "type": "regular",
#           "entities": ["Category", "Category"],
#           "cardinalities": {
#             "Category": "1",
#             "Category": "N"
#           }
#         }
#       ]
#     }

#     Returns:
#         str: The filepath to the generated PNG image file.
#     """
#     try:
#         if not json_input.strip():
#             return "Error: Empty input"
            
#         data = json.loads(json_input)
        
#         if 'entities' not in data:
#             raise ValueError("Missing required field: entities")

#         dot = graphviz.Graph(comment='ER Diagram', engine='dot')
#         dot.attr(rankdir='TB', bgcolor='white', pad='0.5')
#         dot.attr('node', fontname='Arial', fontsize='10', color='black')
#         dot.attr('edge', fontname='Arial', fontsize='9', color='black')
        
#         entities = data.get('entities', [])
#         relationships = data.get('relationships', [])
        
#         for entity in entities:
#             entity_name = entity.get('name')
#             entity_type = entity.get('type', 'strong')
#             attributes = entity.get('attributes', [])
            
#             if not entity_name:
#                 continue
            
#             label_parts = [entity_name]
            
#             if attributes:
#                 attr_lines = []
#                 for attr in attributes:
#                     attr_name = attr.get('name', '')
#                     attr_type = attr.get('type', 'regular')
#                     is_multivalued = attr.get('multivalued', False)
#                     is_derived = attr.get('derived', False)
#                     is_composite = attr.get('composite', False)
                    
#                     if attr_type == 'primary_key':
#                         if is_multivalued:
#                             attr_display = f"PK: {{ {attr_name} }}"
#                         else:
#                             attr_display = f"PK: {attr_name}"
#                     elif attr_type == 'foreign_key':
#                         attr_display = f"FK: {attr_name}"
#                     else:
#                         attr_display = attr_name
#                         if is_derived:
#                             attr_display = f"/ {attr_display} /"
#                         if is_multivalued:
#                             attr_display = f"{{ {attr_display} }}"
#                         if is_composite:
#                             attr_display = f"( {attr_display} )"
                    
#                     attr_lines.append(attr_display)
                
#                 if attr_lines:
#                     label_parts.extend(attr_lines)
            
#             label = "\\n".join(label_parts)
            
#             if entity_type == 'weak':
#                 style = 'filled'
#                 fillcolor = '#e8e8e8'
#                 penwidth = '2'
#             else:
#                 style = 'filled'
#                 fillcolor = '#d0d0d0'
#                 penwidth = '1'
            
#             dot.node(entity_name, label, shape='box', style=style, fillcolor=fillcolor, penwidth=penwidth)
        
#         for relationship in relationships:
#             rel_name = relationship.get('name')
#             rel_type = relationship.get('type', 'regular')
#             entities_involved = relationship.get('entities', [])
#             cardinalities = relationship.get('cardinalities', {})
            
#             if not rel_name or len(entities_involved) < 2:
#                 continue
            
#             if rel_type == 'identifying':
#                 style = 'filled'
#                 fillcolor = '#c0c0c0'
#                 penwidth = '2'
#             else:
#                 style = 'filled'
#                 fillcolor = '#c0c0c0'
#                 penwidth = '1'
            
#             dot.node(rel_name, rel_name, shape='diamond', style=style, fillcolor=fillcolor, penwidth=penwidth)
            
#             for entity in entities_involved:
#                 cardinality = cardinalities.get(entity, '1')
#                 dot.edge(entity, rel_name, label=cardinality)

#         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
#             dot.render(tmp.name, format=output_format, cleanup=True)
#             return f"{tmp.name}.{output_format}"

#     except json.JSONDecodeError:
#         return "Error: Invalid JSON format"
#     except Exception as e:
#         return f"Error: {str(e)}"

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
    """
    Generates an entity relationship diagram from JSON input.

    Args:
        json_input (str): A JSON string describing the entity relationship diagram structure.
                          It must follow the Expected JSON Format Example below.

    Expected JSON Format Example:
    {
      "entities": [
        {
          "name": "EMPLEADO",
          "type": "strong",
          "attributes": [
            {"name": "num_matricula", "type": "primary_key"},
            {"name": "nombre", "type": "regular"},
            {"name": "direccion", "type": "regular"}
          ]
        },
        {
          "name": "DEPARTAMENTO", 
          "type": "strong",
          "attributes": [
            {"name": "cod_depto", "type": "primary_key"},
            {"name": "nombre_depto", "type": "regular"}
          ]
        },
        {
          "name": "TECNICO",
          "type": "strong",
          "attributes": [
            {"name": "nivel", "type": "regular"}
          ]
        },
        {
          "name": "ADMINISTRATIVO",
          "type": "strong", 
          "attributes": []
        },
        {
          "name": "PROYECTO",
          "type": "strong",
          "attributes": [
            {"name": "cod_proyecto", "type": "primary_key"},
            {"name": "nombre_proyecto", "type": "regular"},
            {"name": "cliente", "type": "regular"},
            {"name": "fecha_inicio", "type": "regular"},
            {"name": "fecha_fin", "type": "regular"}
          ]
        }
      ],
      "relationships": [
        {
          "name": "Pertenece",
          "type": "regular",
          "entities": ["EMPLEADO", "DEPARTAMENTO"],
          "cardinalities": {
            "EMPLEADO": "N",
            "DEPARTAMENTO": "1"
          },
          "attributes": [
            {"name": "fecha_asignacion", "type": "regular"},
            {"name": "fecha_cese", "type": "regular"}
          ]
        },
        {
          "name": "Es_un",
          "type": "isa",
          "parent": "EMPLEADO",
          "children": ["TECNICO", "ADMINISTRATIVO"]
        },
        {
          "name": "Trabaja_en",
          "type": "regular",
          "entities": ["TECNICO", "PROYECTO"],
          "cardinalities": {
            "TECNICO": "M",
            "PROYECTO": "N"
          }
        }
      ]
    }

    Returns:
        str: The filepath to the generated PNG image file.
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'entities' not in data:
            raise ValueError("Missing required field: entities")

        # Crear grafo con estilo mejorado
        dot = graphviz.Graph(comment='ER Diagram', engine='dot')
        dot.attr(
            rankdir='TB', 
            bgcolor='white', 
            pad='0.8',
            nodesep='1.2', 
            ranksep='1.8',
            fontname='Arial'
        )
        
        # Configuración base de nodos y aristas
        dot.attr('node', fontname='Arial', fontsize='11', color='#404040')
        dot.attr('edge', fontname='Arial', fontsize='10', color='#404040')
        
        entities = data.get('entities', [])
        relationships = data.get('relationships', [])
        
        # Crear entidades con estilo profesional
        for entity in entities:
            entity_name = entity.get('name')
            entity_type = entity.get('type', 'strong')
            
            if not entity_name:
                continue
            
            if entity_type == 'weak':
                # Entidad débil: doble rectángulo (simulado con tabla HTML)
                label = f'<<TABLE BORDER="3" CELLBORDER="1" CELLSPACING="0" CELLPADDING="8" BGCOLOR="#f0f0f0"><TR><TD><B>{entity_name}</B></TD></TR></TABLE>>'
                dot.node(entity_name, label, shape='plaintext')
            else:
                # Entidad fuerte: rectángulo simple con estilo
                dot.node(
                    entity_name, 
                    entity_name, 
                    shape='box', 
                    style='filled,rounded', 
                    fillcolor='#e8e8e8',
                    color='#404040',
                    penwidth='2',
                    fontsize='12',
                    fontcolor='#000000'
                )
        
        # Crear atributos para cada entidad
        for entity in entities:
            entity_name = entity.get('name')
            attributes = entity.get('attributes', [])
            
            for i, attr in enumerate(attributes):
                attr_name = attr.get('name', '')
                attr_type = attr.get('type', 'regular')
                
                attr_id = f"{entity_name}_attr_{i}"
                
                if attr_type == 'primary_key':
                    # Clave primaria: elipse con texto subrayado
                    dot.node(
                        attr_id, 
                        f'<U><B>{attr_name}</B></U>', 
                        shape='ellipse', 
                        style='filled',
                        fillcolor='#d0d0d0',
                        color='#404040',
                        penwidth='2'
                    )
                elif attr_type == 'partial_key':
                    # Clave parcial: elipse con línea punteada
                    dot.node(
                        attr_id, 
                        f'<U>{attr_name}</U>', 
                        shape='ellipse', 
                        style='filled,dashed',
                        fillcolor='#d0d0d0',
                        color='#404040',
                        penwidth='2'
                    )
                elif attr_type == 'multivalued':
                    # Atributo multivaluado: doble elipse (tabla HTML)
                    label = f'<<TABLE BORDER="2" CELLBORDER="0" CELLSPACING="0" CELLPADDING="4" BGCOLOR="#d8d8d8"><TR><TD>{attr_name}</TD></TR></TABLE>>'
                    dot.node(attr_id, label, shape='plaintext')
                elif attr_type == 'derived':
                    # Atributo derivado: elipse con línea punteada
                    dot.node(
                        attr_id, 
                        f'/{attr_name}/', 
                        shape='ellipse', 
                        style='filled,dashed',
                        fillcolor='#d8d8d8',
                        color='#404040'
                    )
                elif attr_type == 'composite':
                    # Atributo compuesto
                    dot.node(
                        attr_id, 
                        attr_name, 
                        shape='ellipse', 
                        style='filled',
                        fillcolor='#d8d8d8',
                        color='#404040'
                    )
                else:
                    # Atributo regular
                    dot.node(
                        attr_id, 
                        attr_name, 
                        shape='ellipse', 
                        style='filled',
                        fillcolor='#d8d8d8',
                        color='#404040'
                    )
                
                # Conectar atributo con entidad
                dot.edge(entity_name, attr_id, color='#404040', penwidth='1')
        
        # Crear relaciones
        for relationship in relationships:
            rel_name = relationship.get('name')
            rel_type = relationship.get('type', 'regular')
            
            if rel_type == 'isa':
                # Relación ISA (jerarquía)
                parent = relationship.get('parent')
                children = relationship.get('children', [])
                
                if parent and children:
                    # Crear nodo ISA como triángulo
                    isa_id = f"isa_{rel_name}"
                    dot.node(
                        isa_id, 
                        'ISA', 
                        shape='triangle', 
                        style='filled',
                        fillcolor='#c0c0c0',
                        color='#404040',
                        penwidth='2'
                    )
                    
                    # Conectar padre con ISA
                    dot.edge(parent, isa_id, color='#404040', penwidth='2')
                    
                    # Conectar ISA con hijos
                    for child in children:
                        dot.edge(isa_id, child, color='#404040', penwidth='2')
                
            else:
                # Relación regular
                entities_involved = relationship.get('entities', [])
                cardinalities = relationship.get('cardinalities', {})
                rel_attributes = relationship.get('attributes', [])
                
                if not rel_name or len(entities_involved) < 2:
                    continue
                
                rel_type_style = relationship.get('type', 'regular')
                
                if rel_type_style == 'identifying':
                    # Relación identificadora: doble rombo
                    label = f'<<TABLE BORDER="3" CELLBORDER="0" CELLSPACING="0" CELLPADDING="6" BGCOLOR="#c8c8c8"><TR><TD><B>{rel_name}</B></TD></TR></TABLE>>'
                    dot.node(rel_name, label, shape='plaintext')
                else:
                    # Relación regular: rombo simple
                    dot.node(
                        rel_name, 
                        rel_name, 
                        shape='diamond', 
                        style='filled',
                        fillcolor='#c8c8c8',
                        color='#404040',
                        penwidth='2'
                    )
                
                # Crear atributos de la relación
                for i, attr in enumerate(rel_attributes):
                    attr_name = attr.get('name', '')
                    attr_id = f"{rel_name}_attr_{i}"
                    
                    dot.node(
                        attr_id, 
                        attr_name, 
                        shape='ellipse', 
                        style='filled',
                        fillcolor='#d8d8d8',
                        color='#404040'
                    )
                    dot.edge(rel_name, attr_id, color='#404040')
                
                # Conectar entidades con la relación
                for entity in entities_involved:
                    cardinality = cardinalities.get(entity, '1')
                    
                    if len(entities_involved) == 2 and entities_involved[0] == entities_involved[1]:
                        # Relación reflexiva
                        dot.edge(
                            entity, 
                            rel_name, 
                            label=cardinality, 
                            color='#404040',
                            penwidth='2',
                            fontsize='10',
                            fontcolor='#000000'
                        )
                    else:
                        # Relación normal
                        dot.edge(
                            entity, 
                            rel_name, 
                            label=cardinality, 
                            color='#404040',
                            penwidth='2',
                            fontsize='10',
                            fontcolor='#000000'
                        )

        # Renderizar el diagrama
        with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
            dot.render(tmp.name, format=output_format, cleanup=True)
            return f"{tmp.name}.{output_format}"

    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"