import os
import json
import torch
import spaces
import gradio as gr
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from huggingface_hub import HfApi
from datetime import datetime

MODEL_ID = os.environ.get("MODEL_ID")
DATASET_REPO = os.environ.get("DATASET_REPO") 
DESCRIPTION = os.environ.get("DESCRIPTION")
PROMPT = os.environ.get("PROMPT")

ON_LOAD="""
async()=>{
    alert("Before using the service, users must agree to the following terms:\\n\\nPlease note that the model presented here is an experimental tool that is still being developed and improved.\\n\\nMeasures have been taken during the model creation process to minimizing the risk of generating vulgar, prohibited or inappropriate content. However, in rare cases, unwanted content may be generated. If you encounter any content that is deemed inappropriate or violates our policies, please contact us to report it. Your information will enable us to take further steps to improve and develop the model to make it safe and user-friendly.\\n\\nYou must not use the model for illegal, harmful, violent, racist or sexual purposes. Please do not send any private information. The website collects user dialogue data and reserves the right to distribute it under the Creative Commons Attribution (CC-BY) or similar license.");
}
"""

api = HfApi()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


@spaces.GPU()
def generate(instruction, stop_tokens, repetition_penalty, max_new_tokens):
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    input_ids = tokenizer.apply_chat_template(instruction, return_tensors="pt", add_generation_prompt=True,).to(device)

    if input_ids.shape[1] > 4096:
        input_ids = input_ids[:, -4096:]

    generate_kwargs = dict(
        input_ids = input_ids,
        streamer=streamer,
        do_sample=False,
        max_new_tokens=max_new_tokens,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()
    outputs = []
    for new_token in streamer:
        if new_token in stop_tokens:
            break
        outputs.append(new_token.replace("<|im_end|>", ""))
        yield "".join(outputs)


def predict(message, history):
    system_prompt = PROMPT
    repetition_penalty = 1.1
    max_new_tokens = 1024
    stop_tokens = ["<|endoftext|>", "<|im_end|>"]
    conversation = []
    conversation.append(
        {
            "role": "system",
            "content": system_prompt,
        }
    )
    for user, assistant in history:
        conversation.extend([{"role": "user", "content": user}, 
                             {"role": "assistant", "content": assistant}])
    conversation.append({"role": "user", "content": message})
    print(conversation)
    
    for output_text in generate(conversation, stop_tokens, repetition_penalty, max_new_tokens):
        if output_text in stop_tokens:
            break
        yield output_text

    hfapi = HfApi()
    day=datetime.now().strftime("%Y-%m-%d")
    timestamp=datetime.now().timestamp()
    dd={
        'message': message, 
        'history': history, 
        'system_prompt':system_prompt, 
        'max_new_tokens':max_new_tokens, 
        'repetition_penalty':repetition_penalty, 
        'instruction':conversation, 
        'output':output_text,
        'precision': 'auto '+str(model.dtype),
    }
    hfapi.upload_file(
        path_or_fileobj=json.dumps(dd, indent=2, ensure_ascii=False).encode('utf-8'),
        path_in_repo=f"{day}/{timestamp}.json",
        repo_id=DATASET_REPO,
        repo_type="dataset",
        commit_message=f"X",
        run_as_future=True
    )

def vote(chatbot, data: gr.LikeData):
    day=datetime.now().strftime("%Y-%m-%d")
    timestamp=datetime.now().timestamp()
    api.upload_file(
        path_or_fileobj=json.dumps({"history":chatbot, 'index': data.index, 'liked': data.liked}, indent=2, ensure_ascii=False).encode('utf-8'),
        path_in_repo=f"liked/{day}/{timestamp}.json",
        repo_id=DATASET_REPO,
        repo_type="dataset",
        commit_message=f"L",
        run_as_future=True
    )
        
with gr.Blocks(js=ON_LOAD) as demo:
    chatbot = gr.Chatbot(label="Chatbot", likeable=True, render=False)
    chatbot.like(vote, [chatbot], None)
    gr.ChatInterface(
        predict,
        chatbot=chatbot,
        title="MKLLM-7B-Instruct",
        description=DESCRIPTION,
        examples=[
            ["Кој си ти?"],
            ["Колку е 3+6/3-1?"],
            ["Како би го населиле Марс? Биди краток."],
            ["Напиши ми pyhon функција која прима низа броеви и проверува кои броеви се деливи со 7, оние што се деливи ги сместуваме во нова листа која ја враќаме назад."]
        ],
        cache_examples=False
    )
    
if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    tokenizer.eos_token = "<|im_end|>"
    tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
                                                 MODEL_ID,
                                                 device_map=device,
                                                 torch_dtype='auto',
                                                )
    
    demo.queue(max_size=50).launch()