import os
import torch
from threading import Thread
from typing import List, Optional, Tuple, Dict
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import spaces
from pathlib import Path
from huggingface_hub import CommitScheduler
import uuid
import json


# Constants
SYSTEM_PROMPT = """You are a helpful friendly assistant Falcon3 from TII, try to follow instructions as much as possible. Be accurate and brief. The Technology Innovation Institute (TII) is a leading global research center dedicated to pushing the frontiers of knowledge. Our teams of scientists, researchers and engineers work in an open, flexible and agile environment to deliver discovery science and transformative technologies.\nWe are part of the Abu Dhabi Government's Advanced Technology Research Council (ATRC), which oversees technology research in the emirate."""
device = "cuda" if torch.cuda.is_available() else "cpu"
TITLE = "<h1><center>Falcon3 Instruct Playground</center></h1>"
SUB_TITLE = "<h2><center>Try out also <a href='https://chat.falconllm.tii.ae/'>our demo</a> powered by <a href='https://www.openinnovation.ai/'>OpenInnovation AI</a> </center></h2>"

# Custom CSS with dark theme
CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}

h3 {
    text-align: center;
/* Fix for chat container */
.chat-container {
    height: 500px !important;
    overflow-y: auto !important;
    flex-direction: column !important;
}

.messages-container {
    flex-grow: 1 !important;
    overflow-y: auto !important;
    padding-right: 10px !important;
}

.contain {
    height: 100% !important;
}

.model-selector {
    margin: 1em auto !important;
    max-width: 500px !important;
    background: #1f2937 !important;
    padding: 1em !important;
    border-radius: 8px !important;
}

/* Style for radio group container */
.radio-group {
    background: #1f2937 !important;
    padding: 1em !important;
    border-radius: 8px !important;
    gap: 10px !important;
}

/* Style for radio options */
.radio-group label {
    background: #374151 !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 8px 16px !important;
    color: white !important;
}

/* Selected radio option */
.radio-group label[data-selected="true"] {
    background: #6366f1 !important;
    color: white !important;
}

button {
    border-radius: 8px !important;
}
"""

# Global variables to store loaded models and tokenizers
loaded_models = {}
loaded_tokenizers = {}

def load_model(model_size: str):
    """Load model and tokenizer based on selected size"""
    if model_size not in loaded_models:
        model_map = {
            "1B": "tiiuae/Falcon3-1B-Instruct",
            "3B": "tiiuae/Falcon3-3B-Instruct",
            "7B": "tiiuae/Falcon3-7B-Instruct",
            "10B": "tiiuae/Falcon3-10B-Instruct",
            
        }
        
        model_path = model_map[model_size]
        print(f"Loading model: {model_path}")
        
        loaded_models[model_size] = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
        ).to(device)
        loaded_tokenizers[model_size] = AutoTokenizer.from_pretrained(model_path)
    
    return loaded_models[model_size], loaded_tokenizers[model_size]

#load_model("10B")
#load_model("3B")
#load_model("1B")
load_model("7B")


logs_id = os.getenv("LOGS_ID")
logs_token = os.getenv("HF_LOGS_TOKEN")

logs_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
logs_folder = logs_file.parent

scheduler = CommitScheduler(
    repo_id=logs_id,
    repo_type="dataset",
    folder_path=logs_folder,
    path_in_repo="data",
    every=5,
    token=logs_token,
    private=True,
)

@spaces.GPU
def stream_chat(
    message: str,
    history: list,
    model_size: str,
    temperature: float = 0.3,
    max_new_tokens: int = 1024,
    top_p: float = 1.0,
    top_k: int = 20,
    repetition_penalty: float = 1.2,
):
    model, tokenizer = load_model(model_size)
    
    # Create new history list with current message
    new_history = history + [[message, ""]]
    
    conversation = []
    # Only include previous messages in the conversation
    for prompt, answer in history:
        conversation.extend([
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": answer},
        ])
    
    conversation.append({"role": "user", "content": message})
    
    input_text = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(tokenizer, timeout=40.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=inputs,
        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=repetition_penalty,
        streamer=streamer,
        pad_token_id=tokenizer.pad_token_id,
    )
    
    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
    
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        buffer = buffer.replace("\nUser", "")
        buffer = buffer.replace("\nSystem", "")
        new_history[-1][1] = buffer
        yield new_history

    with scheduler.lock:
        with logs_file.open("a") as f:
            f.write(json.dumps({"input": input_text.replace(SYSTEM_PROMPT, ""), "output": buffer.replace(SYSTEM_PROMPT, ""), "model": model_size}))
            f.write("\n")

def clear_input():
    return ""

def add_message(message: str, history: list):
    if message.strip() != "":
        history = history + [[message, ""]]
    return history

def clear_session() -> Tuple[str, List]:
    return '', []

def choose_model(radio: str) -> Tuple[gr.Markdown, gr.Chatbot, str]:
    mark_ = gr.Markdown(value=f"<center><font size=8>Falcon3 {radio} ChatBot</center>")
    chatbot = gr.Chatbot(label=f'Falcon3-{radio}-Instruct', value=[])
    return mark_, chatbot, ""

def main():
    with gr.Blocks(css=CSS, theme="soft") as demo:
        gr.HTML(TITLE)
        gr.HTML(SUB_TITLE)
        gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
        with gr.Row():
            options_models = ["1B", "3B", "7B", "10B"]
            radio = gr.Radio(
                choices=options_models,
                label="Model Size:",
                value="7B",
                elem_classes="radio-group"
            )
        
        with gr.Row():
            with gr.Accordion(label="Chat Interface", open=True):
                mark_ = gr.Markdown("""<center><font size=8>Falcon3 7B ChatBot </center>""")
                
                chatbot = gr.Chatbot(
                    label='Falcon3-7B-Instruct',
                    height=500,
                    container=True,
                    elem_classes=["chat-container"]
                )
                
                with gr.Accordion(label="⚙️ Parameters", open=False):
                    temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.3, label="Temperature")
                    max_new_tokens = gr.Slider(minimum=128, maximum=32768, step=128, value=1024, label="Max new tokens")
                    top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="Top-p")
                    top_k = gr.Slider(minimum=1, maximum=100, step=1, value=20, label="Top-k")
                    repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, step=0.1, value=1.2, label="Repetition penalty")
                
                textbox = gr.Textbox(lines=1, label='Input')
                
                with gr.Row():
                    clear_history = gr.Button("🧹 Clear History")
                    submit = gr.Button("🚀 Send")
                
                # Chain of events for submit button
                    submit_event = submit.click(
                        fn=add_message,
                        inputs=[textbox, chatbot],
                        outputs=chatbot,
                        queue=False
                    ).then(
                        fn=clear_input,
                        outputs=textbox,
                        queue=False
                    ).then(
                        fn=stream_chat,
                        inputs=[textbox, chatbot, radio, 
                               temperature, max_new_tokens, top_p, top_k, repetition_penalty],
                        outputs=chatbot,
                        show_progress=True
                    )

                    # Chain of events for enter key
                    enter_event = textbox.submit(
                        fn=add_message,
                        inputs=[textbox, chatbot],
                        outputs=chatbot,
                        queue=False
                    ).then(
                        fn=clear_input,
                        outputs=textbox,
                        queue=False
                    ).then(
                        fn=stream_chat,
                        inputs=[textbox, chatbot, radio,
                               temperature, max_new_tokens, top_p, top_k, repetition_penalty],
                        outputs=chatbot,
                        show_progress=True
                    )
                    
                    clear_history.click(fn=clear_session,
                                      outputs=[textbox, chatbot])
        
        radio.change(choose_model,
                    inputs=[radio],
                    outputs=[mark_, chatbot, textbox])
        
    
    #demo.queue(api_open=False, default_concurrency_limit=40)
    demo.launch()

if __name__ == "__main__":
    main()