import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Define the system prompt for Editly AI
SYSTEM_PROMPT = "You are Editly AI, an English-speaking assistant specialized in editing and improving text. Communicate with users to enhance their writing. You will answer queries about editing only."

def get_message_tokens(model, role, content):
    content = f"{role}\n{content}\n</s>"
    content = content.encode("utf-8")
    return model.tokenize(content, special=True)

def get_system_tokens(model):
    system_message = {"role": "system", "content": SYSTEM_PROMPT}
    return get_message_tokens(model, **system_message)

def load_model(
    directory: str = ".",
    model_name: str = "editlyai.gguf",
    repo_id: str = "emizemani/editlyai-gguf"
):
    final_model_path = os.path.join(directory, model_name)
    
    print("Downloading model...")
    if not os.path.exists(final_model_path):
        hf_hub_download(
            repo_id=repo_id,
            filename=model_name,
            local_dir=directory,
            use_auth_token=os.getenv('HF_TOKEN')  # Ensure you have the HF_TOKEN set up in your environment
        )
    print("Model downloaded and ready to use.")
    
    model = Llama(
        model_path=final_model_path,
        n_ctx=1024
    )
    
    print("Model loaded!")
    return model

# Initialize the model
MODEL = load_model()

def user(message, history):
    new_history = history + [[message, None]]
    return "", new_history

def bot(history, system_prompt):
    model = MODEL
    tokens = get_system_tokens(model)[:]

    for user_message, bot_message in history[:-1]:
        message_tokens = get_message_tokens(model=model, role="user", content=user_message)
        tokens.extend(message_tokens)
        if bot_message:
            message_tokens = get_message_tokens(model=model, role="bot", content=bot_message)
            tokens.extend(message_tokens)

    last_user_message = history[-1][0]
    message_tokens = get_message_tokens(model=model, role="user", content=last_user_message)
    tokens.extend(message_tokens)

    role_tokens = model.tokenize("bot\n".encode("utf-8"), special=True)
    tokens.extend(role_tokens)
    generator = model.generate(tokens)

    partial_text = ""
    for token in generator:
        if token == model.token_eos():
            break
        partial_text += model.detokenize([token]).decode("utf-8", "ignore")
        history[-1][1] = partial_text
        yield history

# Gradio interface setup
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"<h1>Editly AI - Text Editing Assistant</h1>")
    with gr.Row():
        system_prompt_box = gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt", interactive=False)
        chatbot = gr.Chatbot(label="Conversation")
    with gr.Row():
        msg = gr.Textbox(label="Send a message")
        submit = gr.Button("Send")

    submit.click(
        fn=user,
        inputs=[msg, chatbot],
        outputs=[msg, chatbot]
    ).success(
        fn=bot,
        inputs=[chatbot, system_prompt_box],
        outputs=chatbot
    )

    demo.launch(show_error=True, share=True)