import random

import gradio as gr
import networkx as nx

from lib.graph_extract import triplextract, parse_triples
from lib.visualize import create_graph, create_bokeh_plot, create_plotly_plot
from lib.samples import snippets

WORD_LIMIT = 300

def process_text(text, entity_types, predicates, layout_type, visualization_type):
    if not text:
        return None, None, "Please enter some text."

    words = text.split()
    if len(words) > WORD_LIMIT:
        return None, None, f"Please limit your input to {WORD_LIMIT} words. Current word count: {len(words)}"

    entity_types = [et.strip() for et in entity_types.split(",") if et.strip()]
    predicates = [p.strip() for p in predicates.split(",") if p.strip()]

    if not entity_types:
        return None, None, "Please enter at least one entity type."
    if not predicates:
        return None, None, "Please enter at least one predicate."

    try:
        prediction = triplextract(text, entity_types, predicates)
        if prediction.startswith("Error"):
            return None, None, prediction

        entities, relationships = parse_triples(prediction)

        if not entities and not relationships:
            return None, None, "No entities or relationships found. Try different text or check your input."

        G = create_graph(entities, relationships)

        if visualization_type == 'Bokeh':
            fig = create_bokeh_plot(G, layout_type)
        else:
            fig = create_plotly_plot(G, layout_type)

        output_text = f"Entities: {entities}\nRelationships: {relationships}\n\nRaw output:\n{prediction}"
        return G, fig, output_text
    except Exception as e:
        print(f"Error in process_text: {str(e)}")
        return None, None, f"An error occurred: {str(e)}"

def update_graph(G, layout_type, visualization_type):
    if G is None:
        return None, "Please process text first."
    
    try:
        if visualization_type == 'Bokeh':
            fig = create_bokeh_plot(G, layout_type)
        else:
            fig = create_plotly_plot(G, layout_type)
        return fig, ""
    except Exception as e:
        print(f"Error in update_graph: {e}")
        return None, f"An error occurred while updating the graph: {str(e)}"

def update_inputs(sample_name):
    sample = snippets[sample_name]
    return sample.text_input, sample.entity_types, sample.predicates

with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# Knowledge Graph Extractor")
    
    default_sample_name = random.choice(list(snippets.keys()))
    default_sample = snippets[default_sample_name]
    
    with gr.Row():
        with gr.Column(scale=1):
            sample_dropdown = gr.Dropdown(choices=list(snippets.keys()), label="Select Sample", value=default_sample_name)
            input_text = gr.Textbox(label="Input Text", lines=5, value=default_sample.text_input)
            entity_types = gr.Textbox(label="Entity Types", value=default_sample.entity_types)
            predicates = gr.Textbox(label="Predicates", value=default_sample.predicates)
            layout_type = gr.Dropdown(choices=['spring', 'fruchterman_reingold', 'circular', 'random', 'spectral', 'shell'], 
                                      label="Layout Type", value='spring')
            visualization_type = gr.Radio(choices=['Bokeh', 'Plotly'], label="Visualization Type", value='Bokeh')
            process_btn = gr.Button("Process Text")
        with gr.Column(scale=2):
            output_graph = gr.Plot(label="Knowledge Graph")
            error_message = gr.Textbox(label="Textual Output")

    graph_state = gr.State(None)

    def process_and_update(text, entity_types, predicates, layout_type, visualization_type):
        G, fig, output = process_text(text, entity_types, predicates, layout_type, visualization_type)
        return G, fig, output

    def update_graph_wrapper(G, layout_type, visualization_type):
        if G is not None:
            fig, _ = update_graph(G, layout_type, visualization_type)
            return fig

    sample_dropdown.change(update_inputs, inputs=[sample_dropdown], outputs=[input_text, entity_types, predicates])

    process_btn.click(process_and_update,
                      inputs=[input_text, entity_types, predicates, layout_type, visualization_type],
                      outputs=[graph_state, output_graph, error_message])
    
    layout_type.change(update_graph_wrapper,
                       inputs=[graph_state, layout_type, visualization_type],
                       outputs=[output_graph])
    
    visualization_type.change(update_graph_wrapper,
                              inputs=[graph_state, layout_type, visualization_type],
                              outputs=[output_graph])

if __name__ == "__main__":
    demo.launch(share=True)