# app.py
import gradio as gr
import json
from graphviz import Digraph
import os
from tempfile import NamedTemporaryFile
from sample_data import COMPLEX_SAMPLE_JSON 

def generate_concept_map(json_input: str) -> str:
    """
    Generate concept map from JSON and return as image file
    
    Args:
        json_input (str): JSON describing the concept map structure.
            
            REQUIRED FORMAT EXAMPLE:
            {
                "central_node": "AI",
                "nodes": [
                    {
                        "id": "ml",
                        "label": "Machine Learning",
                        "relationship": "subcategory",
                        "subnodes": [
                            {
                                "id": "dl",
                                "label": "Deep Learning",
                                "relationship": "type",
                                "subnodes": [
                                    {
                                        "id": "cnn",
                                        "label": "CNN",
                                        "relationship": "architecture"
                                    }
                                ]
                            }
                        ]
                    }
                ]
            }
    
    Returns:
        str: Path to generated PNG image file
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'central_node' not in data or 'nodes' not in data:
            raise ValueError("Missing required fields: central_node or nodes")

        # Create graph
        dot = Digraph(
            name='ConceptMap',
            format='png',
            graph_attr={
                'rankdir': 'TB',
                'splines': 'ortho',
                'bgcolor': 'transparent'
            }
        )
        
        # Central node (ellipse)
        dot.node(
            'central',
            data['central_node'],
            shape='ellipse',
            style='filled',
            fillcolor='#2196F3',
            fontcolor='white',
            fontsize='14'
        )
        
        # Process nodes (rectangles)
        for node in data['nodes']:
            node_id = node.get('id')
            label = node.get('label')
            relationship = node.get('relationship')
            
            # Validate node
            if not all([node_id, label, relationship]):
                raise ValueError(f"Invalid node: {node}")
                
            # Create main node (rectangle)
            dot.node(
                node_id,
                label,
                shape='box',
                style='filled',
                fillcolor='#4CAF50',
                fontcolor='white',
                fontsize='12'
            )
            
            # Connect to central node
            dot.edge(
                'central',
                node_id,
                label=relationship,
                color='#9C27B0',
                fontsize='10'
            )
            
            # Helper function to recursively add subnodes and edges
            def add_subnodes(parent_id, subnodes_list, fill_color, font_size, edge_color, edge_font_size):
                for subnode in subnodes_list:
                    sub_id = subnode.get('id')
                    sub_label = subnode.get('label')
                    sub_rel = subnode.get('relationship')
                    
                    if not all([sub_id, sub_label, sub_rel]):
                        raise ValueError(f"Invalid subnode: {subnode}")
                        
                    dot.node(
                        sub_id,
                        sub_label,
                        shape='box',
                        style='filled',
                        fillcolor=fill_color,
                        fontcolor='white',
                        fontsize=str(font_size)
                    )
                    
                    dot.edge(
                        parent_id,
                        sub_id,
                        label=sub_rel,
                        color=edge_color,
                        fontsize=str(edge_font_size)
                    )
                    
                    # Recursively call for deeper levels
                    if 'subnodes' in subnode:
                        # Slightly adjust colors/sizes for deeper levels if desired
                        # For fixed 2 children per parent, you might keep colors consistent per level or vary them.
                        # Here, I'll slightly adjust font size for consistency with depth.
                        add_subnodes(sub_id, subnode['subnodes'], 
                                     '#FFA726' if font_size > 8 else '#FFCC80', # Lighter orange/yellow for deeper levels
                                     font_size - 1 if font_size > 7 else font_size, 
                                     '#E91E63' if edge_font_size > 7 else '#FF5252', # Reddish for deeper edges
                                     edge_font_size - 1 if edge_font_size > 7 else edge_font_size)
            
            # Start processing subnodes from the first level
            add_subnodes(node_id, node.get('subnodes', []), '#FFA726', 10, '#E91E63', 8)


        # Save to temporary file
        with NamedTemporaryFile(delete=False, suffix='.png') as tmp:
            dot.render(tmp.name, format='png', cleanup=True)
            return tmp.name + '.png'

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

if __name__ == "__main__":
    demo = gr.Interface(
        fn=generate_concept_map,
        inputs=gr.Textbox(
            value=COMPLEX_SAMPLE_JSON, # ¡Ahora usa el JSON AI simétrico y complejo!
            placeholder="Paste JSON following the documented format",
            label="Structured JSON Input",
            lines=25
        ),
        outputs=gr.Image(
            label="Generated Concept Map",
            type="filepath",
            show_download_button=True
        ),
        title="Advanced Concept Map Generator (Symmetric AI)",
        description="Create symmetric, multi-level concept maps for AI from properly formatted JSON."
    )
    
    demo.launch(
        mcp_server=True,
        share=False,
        server_port=7860,
        server_name="0.0.0.0"
    )