import gradio as gr
import json
from graphviz import Digraph
import os
from tempfile import NamedTemporaryFile
from sample_data import COMPLEX_SAMPLE_JSON # El JSON de ejemplo se mantiene sin cambios

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")

        dot = Digraph(
            name='ConceptMap',
            format='png',
            graph_attr={
                'rankdir': 'TB',        # Top-to-Bottom
                'splines': 'ortho',     # Straight lines
                'bgcolor': 'white',     # Fondo blanco
                'pad': '0.5',           # Margen alrededor del gráfico
                'size': '8,12!',        # ¡Nuevo! Tamaño deseado (ancho, alto en pulgadas). El '!' es importante.
                'ranksep': '1.0'        # ¡Nuevo! Aumenta la separación vertical entre niveles
                # 'ratio': 'fill'       # Podría usarse con 'size' para forzar a rellenar el espacio
            }
        )
        
        # Base color for the central node
        base_color = '#19191a' # Casi negro

        # Central node (now a rounded box)
        dot.node(
            'central',
            data['central_node'],
            shape='box', # Ahora es un rectángulo
            style='filled,rounded', # Redondeado
            fillcolor=base_color, # Color base
            fontcolor='white',
            fontsize='16' # Un poco más grande para el título
        )
        
        # Helper function to recursively add nodes and edges
        def add_nodes_and_edges(parent_id, nodes_list, current_depth=0):
            # Calculate color for current depth, making it lighter
            lightening_factor = 0.12 # How much lighter each level gets
            
            # Convert base_color hex to RGB
            base_r = int(base_color[1:3], 16)
            base_g = int(base_color[3:5], 16)
            base_b = int(base_color[5:7], 16)

            # Calculate current node color
            current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
            current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
            current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)

            # Clamp values to 255
            current_r = min(255, current_r)
            current_g = min(255, current_g)
            current_b = min(255, current_b)
            
            node_fill_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'

            # Font color: white for dark nodes, black for very light nodes
            font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
            
            # Edge colors can remain constant or change. Let's make them slightly visible.
            edge_color = '#4a4a4a' # Un gris oscuro para las líneas
            font_size = max(9, 14 - (current_depth * 2)) # Adjust font size based on depth
            edge_font_size = max(7, 10 - (current_depth * 1))

            for node in nodes_list:
                node_id = node.get('id')
                label = node.get('label')
                relationship = node.get('relationship')
                
                if not all([node_id, label, relationship]):
                    raise ValueError(f"Invalid node: {node}")
                    
                dot.node(
                    node_id,
                    label,
                    shape='box', # Rectángulo
                    style='filled,rounded', # Redondeado
                    fillcolor=node_fill_color, 
                    fontcolor=font_color,
                    fontsize=str(font_size) 
                )
                
                dot.edge(
                    parent_id,
                    node_id,
                    label=relationship,
                    color=edge_color,
                    fontcolor=edge_color, # Color de la fuente de la arista también gris
                    fontsize=str(edge_font_size) 
                )
                
                if 'subnodes' in node:
                    add_nodes_and_edges(node_id, node['subnodes'], current_depth + 1)

        # Start processing from the top-level nodes connected to the central node
        add_nodes_and_edges('central', data.get('nodes', []), current_depth=1) # Initial depth is 1 for nodes under central

        # 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, 
            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="AI Concept Map (Custom Style)",
        description="Generates an AI concept map with custom rounded boxes, color gradient, and white background."
    )
    
    demo.launch(
        mcp_server=True,
        share=False,
        server_port=7860,
        server_name="0.0.0.0"
    )