import gradio as gr
import numpy as np
from autocorrect_package.autocorrection import Autocorrection

# Initialize spell checker
checker = Autocorrection("autocorrect_package/corpus.txt")

# Shared state
state = {"input_text": "", "last_word": "", "suggestions": []}

# Suggestion generator
def live_autocorrect(input_text):
    state["input_text"] = input_text
    words = input_text.strip().split()

    if not words:
        state["last_word"] = ""
        state["suggestions"] = []
        return input_text, []

    last_word = words[-1].lower()
    state["last_word"] = last_word

    # Already correct
    if last_word in checker.vocabulary:
        state["suggestions"] = []
        return input_text, []

    # Get corrections
    corrections = checker.correct_spelling(last_word)
    corrections = sorted(corrections, key=lambda x: x[1], reverse=True)[:5]
    state["suggestions"] = corrections

    return input_text, [w for w, _ in corrections]

# Replace last word when clicked
def apply_suggestion(evt: gr.SelectData):
    if not evt or not state["input_text"] or not state["last_word"]:
        return state["input_text"]

    selected_word = evt.value.split(" ")[0]
    words = state["input_text"].strip().split()
    if words:
        words[-1] = selected_word
        new_text = " ".join(words)
        state["input_text"] = new_text
        return new_text
    return state["input_text"]

# This function is called on every text change, but only triggers correction if live_mode is True
def trigger_conditionally(input_text, live_mode):
    if live_mode:
        return live_autocorrect(input_text)
    else:
        return gr.update(), gr.update()

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📝 Text Suggestion & Correction")
    gr.Markdown("Toggle 'Live Mode' for real-time suggestions, or use the button.")

    live_toggle = gr.Checkbox(label="Enable Live Mode", value=False)
    input_box = gr.Textbox(label="Type here...", placeholder="Start typing...", lines=2)
    suggest_button = gr.Button("🔍 Suggest Corrections")
    suggestion_list = gr.List(label="Suggestions (click to replace)", interactive=True)

    # When typing, only trigger suggestions if live mode is ON
    input_box.change(fn=trigger_conditionally, inputs=[input_box, live_toggle], outputs=[input_box, suggestion_list])

    # Button trigger (always works)
    suggest_button.click(fn=live_autocorrect, inputs=input_box, outputs=[input_box, suggestion_list])

    # Click to replace word
    suggestion_list.select(fn=apply_suggestion, outputs=input_box)

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