import gradio as gr
import pandas as pd
from utils.ner_helpers import is_llm_model
from typing import Dict, List, Any, Tuple
from tasks.knowledge_graph import build_knowledge_graph, visualize_knowledge_graph_interactive
import base64
from io import BytesIO

def kg_ui():
    """Knowledge Graph UI component"""
    
    # Define models
    KG_MODELS = [
        "gemini-2.0-flash",
        "gpt-4",
        "claude-2",
        "en_core_web_sm",
        "en_core_web_md"
    ]
    DEFAULT_MODEL = "gemini-2.0-flash"
    
    def build_kg(text, model, custom_instructions, interactive=False):
        """Process text for knowledge graph generation"""
        import gradio as gr
        if not text.strip():
            # Trả về các giá trị rỗng cho tất cả các tab
            return (
                "<div style='text-align: center; color: #666; padding: 20px;'>No text provided</div>", 
                pd.DataFrame(), 
                pd.DataFrame(),
                False, True, False, True, False, True
            )
        use_llm = is_llm_model(model)
        result = build_knowledge_graph(
            text=text,
            model_name=model,
            use_llm=use_llm
        )
        entities = result.get("entities", [])
        relations = result.get("relations", [])
        visualization = result.get("visualization")
        # DataFrames
        if entities:
            entities_df = pd.DataFrame(entities)
            entities_df = entities_df.rename(columns={
                "text": "Entity", 
                "label": "Type",
                "start": "Start Position",
                "end": "End Position"
            })
        else:
            entities_df = pd.DataFrame()
        if relations:
            relations_df = pd.DataFrame(relations)
            relations_df = relations_df.rename(columns={
                "subject": "Subject", 
                "relation": "Relation",
                "object": "Object"
            })
        else:
            relations_df = pd.DataFrame()
        # Visualization
        if interactive and entities and relations:
            try:
                interactive_html = visualize_knowledge_graph_interactive(entities, relations)
                visualization_html = f"<div style='width:100%;overflow-x:auto'>{interactive_html}</div>"
                viz_vis = True
                no_viz_vis = False
            except Exception as e:
                visualization_html = f"<div style='color:#d32f2f;padding:20px;'>Error rendering interactive graph: {e}</div>"
                viz_vis = True
                no_viz_vis = False
        elif visualization:
            visualization_html = f"<img src='data:image/png;base64,{visualization}' style='max-width:100%;height:auto;'/>"
            viz_vis = True
            no_viz_vis = False
        else:
            visualization_html = ""
            viz_vis = False
            no_viz_vis = True
        # Visibility flags
        entities_vis = not entities_df.empty
        no_entities_vis = not entities_vis
        relations_vis = not relations_df.empty
        no_relations_vis = not relations_vis
        # Return
        return (
            visualization_html,
            entities_df,
            relations_df,
            viz_vis,
            no_viz_vis,
            entities_vis,
            no_entities_vis,
            relations_vis,
            no_relations_vis
        )
    
    # UI Components
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                label="Input Text",
                lines=8,
                placeholder="Enter text to extract knowledge graph...",
                elem_id="kg-input-text"
            )
            gr.Examples(
                examples=[
                    ["Elon Musk founded SpaceX and Tesla in the United States."],
                    ["Amazon acquired Whole Foods in 2017."]
                ],
                inputs=[input_text],
                label="Examples"
            )
            # Model selection
            model = gr.Dropdown(
                KG_MODELS,
                value=DEFAULT_MODEL,
                label="Model",
                interactive=True,
                elem_id="kg-model-dropdown"
            )
            with gr.Accordion("Advanced Options", open=False, elem_id="kg-advanced-options"):
                custom_instructions = gr.Textbox(
                    label="Custom Instructions",
                    lines=2,
                    placeholder="(Optional) Add specific instructions for knowledge graph generation...",
                    elem_id="kg-custom-instructions"
                )
            btn = gr.Button("Generate Knowledge Graph", elem_id="kg-btn")
        
        with gr.Column(scale=3):
            # Results container with tabs
            with gr.Tabs() as output_tabs:
                with gr.Tab("Graph Visualization", id="kg-viz-tab"):
                    no_viz_html = gr.HTML(
                        "<div style='text-align: center; color: #666; padding: 20px;'>"
                        "Generate a knowledge graph to visualize relationships.</div>",
                        visible=True,
                        elem_id="kg-no-viz"
                    )
                    viz_html = gr.HTML(
                        label="Knowledge Graph Visualization",
                        visible=False,
                        elem_id="kg-viz-html"
                    )
                
                with gr.Tab("Entities", id="kg-entities-tab"):
                    no_entities_html = gr.HTML(
                        "<div style='text-align: center; color: #666; padding: 20px;'>"
                        "No entities found. Try generating a knowledge graph first.</div>",
                        visible=True,
                        elem_id="kg-no-entities"
                    )
                    entities_table = gr.DataFrame(
                        headers=["Entity", "Type", "Start Position", "End Position"],
                        datatype=["str", "str", "number", "number"],
                        visible=False,
                        elem_id="kg-entities-table"
                    )
                
                with gr.Tab("Relationships", id="kg-relations-tab"):
                    no_relations_html = gr.HTML(
                        "<div style='text-align: center; color: #666; padding: 20px;'>"
                        "No relationships found. Try generating a knowledge graph first.</div>",
                        visible=True,
                        elem_id="kg-no-relations"
                    )
                    relations_table = gr.DataFrame(
                        headers=["Subject", "Relation", "Object"],
                        datatype=["str", "str", "str"],
                        visible=False,
                        elem_id="kg-relations-table"
                    )
            
            with gr.Accordion("About Knowledge Graphs", open=False):
                gr.Markdown("""
                ## Knowledge Graphs
                
                Knowledge graphs represent relationships between entities in text as a network. This tool:
                
                - **Extracts entities**: Identifies people, places, organizations, and concepts
                - **Maps relationships**: Shows how entities are connected to each other
                - **Visualizes connections**: Creates an interactive graph you can explore
                
                ### How it works
                
                - **LLM models** can understand complex relationships in text
                - **Traditional models** use pattern matching and syntactic parsing
                
                Knowledge graphs are particularly useful for:
                - Research and analysis
                - Content exploration
                - Understanding complex narratives
                - Discovering hidden connections
                
                Try it with news articles, scientific papers, or story excerpts to see different types of relationships.
                """)
    
    # Toggle for interactive/static visualization
    with gr.Row():
        interactive_toggle = gr.Checkbox(
            label="Interactive Graph (pyvis)",
            value=True,
            elem_id="kg-interactive-toggle"
        )
    # Event handler: use build_kg for all outputs
    def process_and_update_ui(text, model, custom_instructions, interactive):
        return build_kg(text, model, custom_instructions, interactive)
    
    # Wire button to unified handler
    def gradio_output_adapter(visualization_html, entities_df, relations_df, viz_vis, no_viz_vis, entities_vis, no_entities_vis, relations_vis, no_relations_vis):
        return [
            gr.update(value=visualization_html, visible=viz_vis),
            gr.update(value=entities_df, visible=entities_vis),
            gr.update(value=relations_df, visible=relations_vis),
            gr.update(visible=no_viz_vis),
            gr.update(visible=no_entities_vis),
            gr.update(visible=no_relations_vis),
        ]

    btn.click(
        fn=lambda text, model, custom_instructions, interactive: gradio_output_adapter(*build_kg(text, model, custom_instructions, interactive)),
        inputs=[input_text, model, custom_instructions, interactive_toggle],
        outputs=[
            viz_html, entities_table, relations_table,
            no_viz_html, no_entities_html, no_relations_html,
        ]
    )
    return None