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 for the current sentence and last word
state = {"input_text": "", "last_word": "", "suggestions": []}

# Suggestion generator (triggered on text change)
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

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

    # Get corrections
    corrections = checker.correct_spelling(last_word)

    # Sort by confidence descending and keep top 5
    corrections = sorted(corrections, key=lambda x: x[1], reverse=True)[:5]

    state["suggestions"] = corrections

    # suggestion_list = [f"{w} (prob: {round(p, 4)})" for w, p in corrections]
    suggestion_list = [f"{w}" for w, p in corrections]
    return input_text, suggestion_list

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

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📝 Text Suggestion & Correction (Live)")
    gr.Markdown("Type a sentence. Get live suggestions for the last word. Click a suggestion to replace it.")

    input_box = gr.Textbox(label="Type here...", placeholder="Start typing...", lines=2)
    suggestion_list = gr.List(label="Suggestions (click to replace)", interactive=True)

    # Live trigger on text change
    input_box.change(fn=live_autocorrect, inputs=input_box, outputs=[input_box, suggestion_list])

    # Replace last word when clicked
    # suggestion_list.select(fn=apply_suggestion, inputs=suggestion_list, outputs=input_box)
    suggestion_list.select(fn=apply_suggestion, outputs=input_box)

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