import gradio as gr
import requests 
import nltk
import json
import random
from utils import *

nltk.download('punkt')

# Function to handle summarization based on the selected method
def summarize(text: str, method: str) -> str:
    processed_text = ' '.join(text.split("\n"))
    chunks = split_chunk(text)
    tree = MemWalker(chunks)
    tree.build_memory_tree()
    summary = summarize_three_ways(chunks)
    
    #print(processed_text)

    if method == "Truncation":
        return summary['truncated']
    elif method == "Rewrite":
        return summary['rewrite']
    elif method == "Accumulate":
        return summary['accumulate']
    elif method == "Memory Tree":
        return tree.root.summary

example_transcript = get_example()

def clear_input():
    return "", "Truncation"

def load_example(index):
    return example_transcript[index]

def load_example_1():
    return load_example(0)
def load_example_2():
    return load_example(1)
def load_example_3():
    return load_example(2)
def load_example_4():
    return load_example(3)

# Create the interface
with gr.Blocks() as iface:
    text_input = gr.Textbox(lines=5, placeholder="Enter text to summarize here...", label="Input Transcript")
    method_input = gr.Radio(choices=["Truncation", "Rewrite", "Accumulate", "Memory Tree"], label="Summarization Method", value="Truncation")
    output_text = gr.Textbox(label="Summary")

    summarize_button = gr.Button("Summarize")
    clear_button = gr.Button("Clear Input")

    summarize_button.click(summarize, inputs=[text_input, method_input], outputs=output_text)
    clear_button.click(clear_input, outputs=[text_input, method_input])
    #gr.Exa
    #gr.Examples([example_transcript[:100]], inputs=[text_input], label=f"Example", fn=load_example)
    
    example_button_1 = gr.Button(f"Example {1}: {example_transcript[0][:600]}...")
    example_button_1.click(load_example_1, inputs = None, outputs=[text_input])
    
    example_button_2 = gr.Button(f"Example {2}: {example_transcript[1][:600]}...")
    example_button_2.click(load_example_2, inputs = None, outputs=[text_input])
    
    example_button_3 = gr.Button(f"Example {3}: {example_transcript[2][:600]}...")
    example_button_3.click(load_example_3, inputs = None, outputs=[text_input])
    
    example_button_4 = gr.Button(f"Example {4}: {example_transcript[3][:600]}...")
    example_button_4.click(load_example_4, inputs = None, outputs=[text_input])
    

# Launch the interface
iface.launch(share = True, debug = True)