backend/modules/architecture/graph_builder.py

"""
Transform IBM Bob architecture analysis into React Flow graph data.
"""

from __future__ import annotations

NODE_COLORS = {
    "frontend": "#3B82F6",
    "backend": "#8B5CF6",
    "database": "#10B981",
    "api": "#F59E0B",
    "middleware": "#6366F1",
    "utility": "#64748B",
    "auth": "#EF4444",
    "payment": "#F97316",
    "notification": "#06B6D4",
    "code": "#84CC16",
    "service": "#EC4899",
    "queue": "#F97316",
    "cloud": "#14B8A6",
}


def build_workflow_from_analysis(analysis: dict) -> dict:
    """Convert architecture analysis into workflow nodes and edges."""
    workflows = _ensure_list(analysis.get("business_workflows", []))
    
    if not workflows:
        return {"nodes": [], "edges": []}
    
    # Take the first workflow (unified project workflow)
    workflow = workflows[0] if workflows else {}
    steps = _ensure_list(workflow.get("steps", []))
    services_involved = _ensure_list(workflow.get("services_involved", []))
    
    nodes = []
    edges = []
    
    # Create nodes from workflow steps
    for idx, step in enumerate(steps):
        step_text = str(step)
        # Infer node type from step text
        node_type = "process"
        color = "#3B82F6"
        
        if any(word in step_text.lower() for word in ["start", "trigger", "initiate", "begin"]):
            node_type = "start"
            color = "#10B981"
        elif any(word in step_text.lower() for word in ["check", "validate", "verify", "decision"]):
            node_type = "decision"
            color = "#F59E0B"
        elif any(word in step_text.lower() for word in ["database", "store", "save", "persist"]):
            node_type = "database"
            color = "#10B981"
        elif any(word in step_text.lower() for word in ["api", "request", "call", "fetch"]):
            node_type = "api"
            color = "#06B6D4"
        elif any(word in step_text.lower() for word in ["auth", "login", "authenticate"]):
            node_type = "auth"
            color = "#EF4444"
        elif any(word in step_text.lower() for word in ["notify", "alert", "email", "message"]):
            node_type = "notification"
            color = "#8B5CF6"
        
        nodes.append({
            "id": f"workflow_step_{idx}",
            "type": "workflowNode",
            "position": {
                "x": 150 + (idx % 3) * 280,
                "y": 100 + (idx // 3) * 180
            },
            "data": {
                "label": step_text[:40],  # Truncate long labels
                "nodeType": node_type,
                "color": color,
                "description": step_text,
            }
        })
        
        # Create edge to next step
        if idx < len(steps) - 1:
            edges.append({
                "id": f"workflow_edge_{idx}",
                "source": f"workflow_step_{idx}",
                "target": f"workflow_step_{idx + 1}",
                "type": "workflowEdge",
                "animated": True,
            })
    
    return {
        "nodes": nodes,
        "edges": edges,
        "workflow_name": workflow.get("name", "Project Workflow"),
    }


def _ensure_list(value):
    if value is None:
        return []
    if isinstance(value, list):
        return value
    if isinstance(value, tuple) or isinstance(value, set):
        return list(value)
    return [value]


def build_react_flow_graph(analysis: dict) -> dict:
    """Convert architecture analysis to React Flow nodes and edges with improved layout."""
    nodes = []
    edges = []
    services = _ensure_list(analysis.get("core_services", []))
    dependencies = _ensure_list(analysis.get("critical_dependencies", []))
    workflows = _ensure_list(analysis.get("business_workflows", []))

    layer_order = ["frontend", "api", "middleware", "backend", "database", "auth", "utility", "notification"]
    group_order = ["pages", "ui_domain", "client_core", "api_router", "backend_module", "ai", "backend_core", ""]

    buckets: dict[tuple[str, str], list[dict]] = {}
    for service in services:
        layer_type = service.get("type", "utility")
        group = service.get("group") or ""
        buckets.setdefault((layer_type, group), []).append(service)

    y_offset = 80
    layer_spacing = 240
    node_spacing = 220

    for layer_type in layer_order:
        layer_groups = [group for (layer, group) in buckets if layer == layer_type]
        ordered_groups = [group for group in group_order if group in layer_groups]
        ordered_groups.extend([group for group in layer_groups if group not in ordered_groups])

        for group in ordered_groups:
            layer_services = buckets.get((layer_type, group), [])
            if not layer_services:
                continue

            total_width = len(layer_services) * node_spacing
            x_start = max(80, (1280 - total_width) // 2)

            for idx, service in enumerate(layer_services):
                service_id = service.get("id") or f"service_{len(nodes)}"
                nodes.append(
                    {
                        "id": service_id,
                        "type": "serviceNode",
                        "position": {"x": x_start + (idx * node_spacing), "y": y_offset},
                        "data": {
                            "label": service.get("name", service_id),
                            "type": layer_type,
                            "description": service.get("description", ""),
                            "critical": service.get("critical", False),
                            "files": [str(file) for file in _ensure_list(service.get("files", [])) if file],
                            "color": NODE_COLORS.get(layer_type, "#64748B"),
                            "group": group,
                        },
                    }
                )
            y_offset += layer_spacing

    # Build edges with better styling
    known_ids = {node["id"] for node in nodes}
    for index, dep in enumerate(dependencies):
        source = dep.get("from")
        target = dep.get("to")
        if source not in known_ids or target not in known_ids:
            continue
        edge_type = dep.get("type", "imports")
        edges.append({
            "id": f"e{index}",
            "source": source,
            "target": target,
            "type": "default",
            "label": dep.get("description", edge_type)[:30],  # Truncate long labels
            "animated": edge_type in {"api_call", "event"},
            "style": {
                "stroke": "#6366F1" if edge_type == "api_call" else "#64748B",
                "strokeWidth": 2 if dep.get("critical") else 1.5,
            },
            "data": {"edgeType": edge_type},
        })

    return {
        "nodes": nodes,
        "edges": edges,
        "workflows": workflows,
        "mermaid_diagram": analysis.get("mermaid_diagram", ""),
        "metadata": {
            "architecture_type": analysis.get("architecture_type", "Unknown"),
            "summary": analysis.get("summary", ""),
            "entry_points": analysis.get("entry_points", []),
            "tech_decisions": analysis.get("tech_decisions", []),
            "mermaid_diagram": analysis.get("mermaid_diagram", ""),
        },
    }