import os
import gradio as gr
from huggingface_hub import HfApi, login, upload_file

rules = {}

def load_rules():
    global rules
    try:
        with open("rules.txt", "r") as f:
            lines = f.readlines()
            rules = dict(line.strip().split("=>", 1) for line in lines if "=>" in line)
    except Exception:
        rules = {}

def chatbot_response(message, history):
    # Always works, regardless of file size
    for k, v in rules.items():
        if k in message:
            return v
    return "No matching rule. Submit an example!"

def submit_rule(user_input, rule_response, write_token):
    import gradio as gr
    global rules
    filesize_mb = os.path.getsize("rules.txt")/1024/1024 if os.path.exists("rules.txt") else 0
    if filesize_mb > 950:
        raise gr.Error("Dataset is full. No further submissions allowed.")
    rules[user_input] = rule_response
    try:
        with open("rules.txt", "a") as f:
            f.write(f"{user_input}=>{rule_response}\n")
    except Exception as e:
        raise gr.Error(f"Couldn't write rule file: {e}")
    if write_token:
        try:
            login(token=write_token)
            api = HfApi()
            upload_file(
                repo_id="JaySStest/shittyassspace",
                path_or_fileobj="rules.txt",
                path_in_repo="rules.txt",
                repo_type="space",
                token=write_token
            )
        except Exception as e:
            raise gr.Error(f"Failed to update remote dataset: {e}")
    return "Rule submitted!"

load_rules()

with gr.Blocks() as demo:
    chat = gr.ChatInterface(
        fn=chatbot_response,
        type="messages"
    )
    rule_input = gr.Textbox(label="User Example")
    rule_output = gr.Textbox(label="Response")
    token_input = gr.Textbox(label="Write Token", type="password")
    submit = gr.Button("Submit Rule")
    token_help = gr.Button("Don't have a token? Make one.", link="https://huggingface.co/settings/tokens")
    status_output = gr.Textbox(label="Status", interactive=False)
    submit.click(
        submit_rule,
        inputs=[rule_input, rule_output, token_input],
        outputs=status_output
    )

demo.launch()