import os
import json
import subprocess
from threading import Thread

import torch
import spaces
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
from sentence_transformers import SentenceTransformer
from datasets import load_dataset


subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

MODEL_ID = os.environ.get("MODEL_ID")
CHAT_TEMPLATE = os.environ.get("CHAT_TEMPLATE")
MODEL_NAME = MODEL_ID.split("/")[-1]
CONTEXT_LENGTH = int(os.environ.get("CONTEXT_LENGTH"))
COLOR = os.environ.get("COLOR")
EMOJI = os.environ.get("EMOJI")
DESCRIPTION = os.environ.get("DESCRIPTION")
token = os.environ["HF_TOKEN"]

# ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# dataset = load_dataset("gagan3012/edgar-corpus-embeddings")

# data = dataset["train"]
# data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset

def search(query: str, k: int = 3 ):
    """a function that embeds a new query and returns the most probable results"""
    embedded_query = ST.encode(query) # embed new query
    scores, retrieved_examples = data.get_nearest_examples( # retrieve results
        "embeddings", embedded_query, # compare our new embedded query with the dataset embeddings
        k=k # get only top k results
    )
    return scores, retrieved_examples

def format_prompt(prompt,retrieved_documents,k):
    """using the retrieved documents we will prompt the model to generate our responses"""
    PROMPT = f"Question:{prompt}\nContext:"
    for idx in range(k) :
        PROMPT+= f"{retrieved_documents['text'][idx]}\n"
    return PROMPT


# @spaces.GPU()
# def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
#     # Format history with a given chat template
#     system_prompt = """You are FinTral, a financial expert specializing in the nuanced analysis
# of financial statements and a wide array of data-driven financial
# tasks. For each prompt you are given, think step by step. Sometimes, you must extract relevant information to proceed with the
# problem.
# Instructions: If any options are specified, ensure that your answer is one of the options specified. Do not explain why you think the answer is correct.
# """
#     if CHAT_TEMPLATE == "ChatML":
#         stop_tokens = ["<|endoftext|>", "<|im_end|>", "<|end_of_text|>", "<|eot_id|>", "assistant"]
#         instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
#         # scores , retrieved_documents = search(message, k=2)
#         # final_message = format_prompt(message,retrieved_documents,k=2)
#         for human, assistant in history:
#             instruction += '<|im_start|>user\n' + human + '\n<|im_end|>\n<|im_start|>assistant\n' + assistant
#         instruction += '\n<|im_start|>user\n' + message + '\n<|im_end|>\n<|im_start|>assistant\n'
#     elif CHAT_TEMPLATE == "Mistral Instruct":
#         stop_tokens = ["</s>", "[INST]", "[INST] ", "<s>", "[/INST]", "[/INST] "]
#         instruction = '<s>[INST] ' + system_prompt
#         for human, assistant in history:
#             instruction += human + ' [/INST] ' + assistant + '</s>[INST]'
#         instruction += ' ' + message + ' [/INST]'
#     else:
#         messages = [{"role":"system","content":system_prompt},{"role":"user","content":message}]
#         for human, assistant in history:
#             messages.append({"role":"user","content":human},{"role":"assistant","content":assistant})
#     # tell the model to generate
#         instruction = tokenizer.apply_chat_template(
#             messages,
#             add_generation_prompt=True,
#             tokenize=False
#         )
#     print(instruction)
    
#     streamer = TextIteratorStreamer(tokenizer, timeout=90.0, skip_prompt=True, skip_special_tokens=True)
#     enc = tokenizer([instruction], return_tensors="pt", padding=True, truncation=True, max_length=CONTEXT_LENGTH)
#     input_ids, attention_mask = enc.input_ids, enc.attention_mask

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

#     generate_kwargs = dict(
#         {"input_ids": input_ids.to(device), "attention_mask": attention_mask.to(device)},
#         streamer=streamer,
#         do_sample=True,
#         temperature=temperature,
#         max_new_tokens=max_new_tokens,
#         top_k=top_k,
#         repetition_penalty=repetition_penalty,
#         top_p=top_p
#     )
#     t = Thread(target=model.generate, kwargs=generate_kwargs)
#     t.start()
#     outputs = []
#     for new_token in streamer:
#         outputs.append(new_token)
#         if new_token in stop_tokens:
#             break
#         yield "".join(outputs)


@spaces.GPU()
def stream_chat(
    message: str, 
    history: list,
    system_prompt: str,
    temperature: float = 0.8, 
    max_new_tokens: int = 1024, 
    top_p: float = 1.0, 
    top_k: int = 20, 
    penalty: float = 1.2,
):
    print(f'message: {message}')
    print(f'history: {history}')
    system = """You are FinTral, a financial expert specializing in the nuanced analysis
of financial statements and a wide array of data-driven financial
tasks. For each prompt you are given, think step by step. Sometimes, you must extract relevant information to proceed with the
problem.
Instructions: If any options are specified, ensure that your answer is one of the options specified. Do not explain why you think the answer is correct.
"""

    conversation = [
        {"role": "system", "content": system}
    ]
    for prompt, answer in history:
        conversation.extend([
            {"role": "user", "content": prompt}, 
            {"role": "assistant", "content": answer},
        ])

    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
    
    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=input_ids, 
        max_new_tokens = max_new_tokens,
        do_sample = False if temperature == 0 else True,
        top_p = top_p,
        top_k = top_k,
        temperature = temperature,
        repetition_penalty=1.2,
        eos_token_id=[128001,128008,128009],
        streamer=streamer,
    )

    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer

# Load model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type= "nf4"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID,token=token)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    quantization_config=quantization_config,
    attn_implementation="flash_attention_2",
    token=token
)

css = """
.message-row {
    justify-content: space-evenly !important;
}
.message-bubble-border {
    border-radius: 6px !important;
}
.message-buttons-bot, .message-buttons-user {
    right: 10px !important;
    left: auto !important;
    bottom: 2px !important;
}
.dark.message-bubble-border {
    border-color: #15172c !important;
}
.dark.user {
    background: #10132c !important;
}
.dark.assistant.dark, .dark.pending.dark {
    background: #020417 !important;
}
"""

TITLE = "<h1><center>FinTral</center></h1>"

PLACEHOLDER = """
<center>
<p>Hi! How can I help you today?</p>
</center>
"""


CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
"""

# Create Gradio interface
chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)

with gr.Blocks(css=CSS, theme="soft") as demo:
    gr.HTML(TITLE)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
    gr.ChatInterface(
        fn=stream_chat,
        chatbot=chatbot,
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Textbox(
                value="You are a helpful assistant",
                label="System Prompt",
                render=False,
            ),
            gr.Slider(
                minimum=0,
                maximum=1,
                step=0.1,
                value=0.8,
                label="Temperature",
                render=False,
            ),
            gr.Slider(
                minimum=128,
                maximum=8192,
                step=1,
                value=1024,
                label="Max new tokens",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.1,
                value=1.0,
                label="top_p",
                render=False,
            ),
            gr.Slider(
                minimum=1,
                maximum=20,
                step=1,
                value=20,
                label="top_k",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=2.0,
                step=0.1,
                value=1.2,
                label="Repetition penalty",
                render=False,
            ),
        ],
        examples=[
            ["Help me study vocabulary: write a sentence for me to fill in the blank, and I'll try to pick the correct option."],
            ["What are 5 creative things I could do with my kids' art? I don't want to throw them away, but it's also so much clutter."],
            ["Tell me a random fun fact about the Roman Empire."],
            ["Show me a code snippet of a website's sticky header in CSS and JavaScript."],
        ],
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.launch()