generate_data.py

import json
import random
import string
from typing import List, Dict


def generate_random_id(length: int = 8) -> str:
    """Genera un ID aleatorio"""
    return "".join(random.choices(string.ascii_letters + string.digits, k=length))


def generate_nodes(count: int) -> List[Dict]:
    """Genera una cantidad específica de nodos"""
    nodes = []
    node_types = [
        "text_message",
        "quick_reply",
        "card",
        "carousel",
        "image",
        "video",
        "audio",
        "file",
    ]

    for i in range(count):
        # Nodo inicial especial
        if i == 0:
            node = {
                "id": "node-START",
                "data": {"type": "START", "label": "Inicio del flujo"},
                "type": "custom",
                "width": 108,
                "height": 42,
                "position": {
                    "x": random.randint(0, 1000),
                    "y": random.randint(0, 1000),
                },
                "deletable": False,
            }
        else:
            # Algunos nodos especiales para testing
            if i < 10:
                node_id = f"node-{generate_random_id()}"
            else:
                node_id = f"node-{i:06d}"

            node_type = random.choice(node_types)

            node = {
                "id": node_id,
                "data": {
                    "type": node_type,
                    "label": f"Nodo {node_type} #{i}",
                    "message": f"Mensaje del nodo {i}: {generate_random_id(20)}",
                    "saveResponse": random.choice([True, False]),
                },
                "type": "custom",
                "width": random.randint(200, 400),
                "height": random.randint(100, 300),
                "position": {
                    "x": random.randint(0, 5000),
                    "y": random.randint(0, 5000),
                },
                "dragging": False,
                "selected": False,
            }

            # Agregar botones a algunos nodos
            if random.random() < 0.3:  # 30% de probabilidad
                buttons = []
                num_buttons = random.randint(1, 4)
                for b in range(num_buttons):
                    button_types = ["QUICK_REPLY", "POSTBACK", "URL", "PHONE"]
                    buttons.append(
                        {
                            "id": generate_random_id(12),
                            "text": f"Botón {b + 1}",
                            "type": random.choice(button_types),
                        }
                    )
                node["data"]["buttons"] = buttons

        nodes.append(node)

    return nodes


def generate_edges(nodes: List[Dict], density: float = 0.15) -> List[Dict]:
    """
    Genera edges entre nodos con una densidad específica
    density: porcentaje de conexiones posibles que se crearán
    """
    edges = []
    node_ids = [node["id"] for node in nodes]

    # Asegurar que cada nodo tenga al menos una conexión de salida (excepto algunos finales)
    for i, source_id in enumerate(
        node_ids[: -int(len(node_ids) * 0.1)]
    ):  # 90% de nodos tienen salida
        # Número de conexiones de salida por nodo
        num_connections = random.randint(1, min(5, len(node_ids) - i - 1))

        # Seleccionar nodos destino (evitar auto-referencias y ciclos simples)
        possible_targets = node_ids[i + 1 : i + 1 + min(20, len(node_ids) - i - 1)]
        targets = random.sample(
            possible_targets, min(num_connections, len(possible_targets))
        )

        for j, target_id in enumerate(targets):
            edge_types = ["custom-edge", "default", "smoothstep", "step"]
            button_types = ["QUICK_REPLY", "POSTBACK", "MESSAGE", "URL"]

            edge = {
                "id": f"edge_{source_id}_{target_id}_{j}",
                "type": random.choice(edge_types),
                "source": source_id,
                "target": target_id,
                "animated": random.choice([True, False]),
            }

            # Agregar sourceHandle para algunos edges (importante para criterios de búsqueda)
            if random.random() < 0.4:  # 40% tienen sourceHandle
                button_type = random.choice(button_types)
                edge["sourceHandle"] = f"source-{source_id}-{button_type}-{j}"

            # Agregar targetHandle para algunos edges
            if random.random() < 0.3:  # 30% tienen targetHandle
                edge["targetHandle"] = f"target-{target_id}"

            # Agregar label a algunos edges
            if random.random() < 0.2:  # 20% tienen label
                edge["label"] = f"Conexión: {source_id[:8]}→{target_id[:8]}"

            edges.append(edge)

    return edges


def generate_large_dataset(num_nodes: int, edge_density: float = 0.15):
    """Genera un dataset grande y lo guarda en archivos JSON"""

    print(f"🏗️  Generando {num_nodes:,} nodos...")
    nodes = generate_nodes(num_nodes)

    print(f"🔗 Generando edges con densidad {edge_density:.1%}...")
    edges = generate_edges(nodes, edge_density)

    print(f"📁 Guardando {len(nodes):,} nodos y {len(edges):,} edges...")

    # Guardar en archivos JSON
    with open("data/nodes_large.json", "w", encoding="utf-8") as f:
        json.dump(nodes, f, indent=2, ensure_ascii=False)

    with open("data/edges_large.json", "w", encoding="utf-8") as f:
        json.dump(edges, f, indent=2, ensure_ascii=False)

    print(f"✅ Dataset generado:")
    print(f"   • Nodos: {len(nodes):,}")
    print(f"   • Edges: {len(edges):,}")
    print(f"   • Ratio edges/nodos: {len(edges) / len(nodes):.2f}")

    return nodes, edges


if __name__ == "__main__":
    # Generar diferentes tamaños para testing
    sizes = [1000, 5000, 10000]

    for size in sizes:
        print(f"\n{'=' * 60}")
        print(f"Generando dataset de {size:,} nodos...")
        nodes, edges = generate_large_dataset(size)

        # Renombrar archivos por tamaño
        import os

        os.rename("data/nodes_large.json", f"data/nodes_{size}.json")
        os.rename("data/edges_large.json", f"data/edges_{size}.json")