import gradio
from huggingface_hub import InferenceClient

import datetime
import uuid
import json

import re

import os
import sys

history = [
    {"role": "system", "content": ""},
]


tmp_dir = os.getenv('GRADIO_TEMP_DIR')


def generate_uuid():
    _uuid = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + str(uuid.uuid4())
    return _uuid


def write_current_chat_to_file(current_chat):
    os.makedirs(tmp_dir, exist_ok = True)
    with open(f"{tmp_dir}/{current_chat['chat_id']}.json", "w") as f:
        json.dump(current_chat, f, indent = 4)
    
    return gradio.DownloadButton(f"{tmp_dir}/{current_chat['chat_id']}.json")
        

def process_input_message(message_box, current_chat):
    current_chat["chat_history"].append({"role": "user", "content": message_box["text"]})
    return current_chat


def get_text_between_tags(text, start_tag, end_tag):
    pattern = rf'{re.escape(start_tag)}(.*?){re.escape(end_tag)}'
    match = re.search(pattern, text, re.DOTALL)
    return match.group(1) if match else ""


def remove_text_between_tags(text, start_tag, end_tag):
    pattern = rf'{re.escape(start_tag)}.*?{re.escape(end_tag)}'
    return re.sub(pattern, '', text, flags=re.DOTALL)


def call_chatbot(api_token, current_chat, system_message, max_tokens, temperature, top_p, use_thoughts_as_context):
    client = InferenceClient(
        provider = "hf-inference",
        api_key = api_token
    )

    current_chat["chat_history"][0]["content"] = system_message
    current_chat["chat_history"].append({"content": "", "role": "assistant"})
    
    messages = current_chat["chat_history"] if use_thoughts_as_context else [message for message in current_chat["chat_history"] if "metadata" not in message.keys()]
    print(messages)
    
    stream = client.chat.completions.create(
        model = "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", 
        messages = current_chat["chat_history"], 
        max_tokens = max_tokens,
        temperature = temperature,
        top_p = top_p,
        stream = True
    )

    response = ""
    for chunk in stream:
        response = response + chunk.choices[0].delta.content
        current_chat["chat_history"][-1]["content"] = response
        yield current_chat["chat_history"], current_chat
        #yield current_chat["chat_history"], current_chat
    
    current_chat["chat_history"][-1] = {
        "content": get_text_between_tags("<begin>" + response, "<begin>", "</think>").replace("<think>", ""), 
        "role": "assistant", 
        "metadata": {
            "title": "💭 Thoughts:"
        }}
    current_chat["chat_history"].append({"content": get_text_between_tags(response + "</end>", "</think>", "</end>"), "role": "assistant"})
    
    yield current_chat["chat_history"], current_chat
    
    
def reset_chat():
    chatbot = gradio.Chatbot(history, type = "messages")
    message_box = gradio.MultimodalTextbox(
        value = "",
        interactive = True,
        file_count = "multiple",
        placeholder = "Enter message...",
        show_label = False,
        sources = [],
        stop_btn = True,
    )
    
    current_chat_id = generate_uuid()
    current_chat = gradio.JSON(
        {
            "version": "0.1",
            "chat_id": current_chat_id, 
            "chat_history": history
        }
    )
    download_file = gradio.File(label = "Save", value = f"{tmp_dir}/{current_chat_id}.json")
    
    return chatbot, message_box, current_chat, download_file


def reset_parameters():
    system_message = gradio.Textbox(label = "System Message", value = "You are a helpful bot. Be concise with your answers. Do not think with more than 3 lines. Answer in 2 lines. Only answer in English.")
    max_tokens = gradio.Slider(label = "Max Tokens", minimum = 500, maximum = 3000, step = 100, value = 1000)
    temperature = gradio.Slider(label = "Temperature", minimum = 0.1, maximum = 2.0, step = 0.1, value = 0.5)
    top_p = gradio.Slider(label = "Top P", minimum = 0.1, maximum = 1.0, step = 0.1, value = 0.9)
    use_thoughts_as_context = gradio.Checkbox(value = False, label = "Use thoughts as context")
    return  system_message, max_tokens, temperature, top_p, use_thoughts_as_context


def process_token(secret_token):
    try:
        passwords = os.environ.get("PASSWORDS")
        passwords = passwords.split(":")
        
        if secret_token in passwords:
            secret_token = os.environ.get("HF_KEY")
        
        return secret_token
    except:
        return secret_token
    
    
with gradio.Blocks(fill_height = True) as base_app:
    gradio.Markdown("# ChatSeek")
    gradio.Markdown("## ")
    
    with gradio.Row():
        with gradio.Column(scale = 2):
            secret_token = gradio.Textbox(label = "API Key", placeholder = "Enter Password/API Token. The key is never stored.", type = "password")
    
    chatbot = gradio.Chatbot(history, type = "messages")
    message_box = gradio.MultimodalTextbox(
        interactive = True,
        file_count = "multiple",
        placeholder = "Enter message...",
        show_label = False,
        sources = [],
        stop_btn = True,
    )
    
    current_chat_id = generate_uuid()
    
    with gradio.Row(equal_height = True):
        with gradio.Column():
            reset_chat_button = gradio.Button(value = "Start a New Chat")
        with gradio.Column():
            save_chat_button = gradio.DownloadButton(label = "Save", value = f"{tmp_dir}/{current_chat_id}.json")
            
    with gradio.Accordion(label = "Advanced Parameters", open = False):
        system_message = gradio.Textbox(label = "System Message", value = "You are a helpful bot. Be concise with your answers. Do not think with more than 3 lines. Answer in 2 lines. Only answer in English.")
        max_tokens = gradio.Slider(label = "Max Tokens", minimum = 500, maximum = 3000, step = 100, value = 1000)
        temperature = gradio.Slider(label = "Temperature", minimum = 0.1, maximum = 2.0, step = 0.1, value = 0.5)
        top_p = gradio.Slider(label = "Top P", minimum = 0.1, maximum = 1.0, step = 0.1, value = 0.9)
        use_thoughts_as_context = gradio.Checkbox(value = False, label = "Use thoughts as context")
        
        reset_parameters_button = gradio.Button(value = "Reset Parameters")
    
    with gradio.Accordion(label = "Metadata", open = False):
        current_chat = gradio.JSON(
            {
                "version": "0.1",
                "chat_id": current_chat_id, 
                "chat_history": history
            },
            visible = True
        )
    
    secret_token_submit_call = secret_token.submit(process_token, [secret_token], [secret_token])
    
    submit_message_call = message_box.submit(process_input_message, [message_box, current_chat], [current_chat], queue=False).then(write_current_chat_to_file, [current_chat], [save_chat_button])
    clear_message_box_call = submit_message_call.then(lambda: gradio.MultimodalTextbox(value = "", interactive = True) , None, [message_box])
    invoke_chatbot_call = clear_message_box_call.then(call_chatbot, [secret_token, current_chat, system_message, max_tokens, temperature, top_p, use_thoughts_as_context], [chatbot, current_chat]).then(write_current_chat_to_file, [current_chat], [save_chat_button])
    
    reset_chat_button_call = reset_chat_button.click(reset_chat, [], [chatbot, message_box, current_chat])
    reset_parameters_button_call = reset_parameters_button.click(reset_parameters, [], [system_message, max_tokens, temperature, top_p, use_thoughts_as_context])
    
if __name__ == "__main__":
    base_app.launch(
        allowed_paths = [tmp_dir]
    )