import gradio as gr
from collections import defaultdict
import jieba

def get_words():
    with open('./data.txt', 'r', encoding='utf-8') as f:
        data = f.read()
    # Use jieba for word segmentation
    words = jieba.cut(data)

    # Convert the result to a list or output as a string
    words_list = list(words)
    words_list = [word for word in words_list if word.replace(' ', '') not in ['', '\n', ',', '.', '"', '“', '”', '’', '‘', '!', '-', '#', '`', '*']]

    return words_list

class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.is_end_of_word = False
        self.frequency = 0 

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        node = self.root
        for char in word:
            node = node.children[char]
        node.is_end_of_word = True
        node.frequency += 1

    def search(self, prefix):
        node = self.root
        for char in prefix:
            if char not in node.children:
                return []
            node = node.children[char]
        return self._find_words_from_node(node, prefix)

    def _find_words_from_node(self, node, prefix):
        words = []
        if node.is_end_of_word:
            words.append((prefix, node.frequency))
        for char, next_node in node.children.items():
            words.extend(self._find_words_from_node(next_node, prefix + char))
        words.sort(key=lambda x: x[1], reverse=True)
        return words

# Gradio interface function
def autocomplete(input_text):
    if ' ' in input_text:
        input_text = input_text.split(' ')[-1]
    if len(input_text) == 0:
        return ''
    suggestions = trie.search(input_text)
    # Format suggestions as a string with one word per line
    formatted_suggestions = "\n".join(s[0] for s in suggestions)
    # Return the formatted string
    return formatted_suggestions if len(formatted_suggestions) > 0 else input_text

if __name__ == '__main__':
    words = get_words()
    trie = Trie()
    for word in words:
        trie.insert(word)

    # Create Gradio interface
    iface = gr.Interface(
        fn=autocomplete,  # Function to call
        inputs=gr.Textbox(lines=1, placeholder="Type something..."),  # Input textbox
        outputs=gr.Textbox(lines=5, label="Suggestions"),  # Output textbox
        live=True,  # Enable live updates
        title="String Autocompletion Demo",  # Interface title
        description="Type a prefix to see autocomplete suggestions based on word frequency."
    )

    # Launch the interface
    iface.launch()