import os
from threading import Thread
from typing import Iterator, List, Tuple
import json

import gradio as gr
import spaces
import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from flask import Flask, request, jsonify

DESCRIPTION = """\
# Zero GPU Model Comparison Arena
Compare two language models using Hugging Face's Zero GPU initiative.
Select two different models from the dropdowns and see how they perform on the same input.
"""

MAX_MAX_NEW_TOKENS = 1024
DEFAULT_MAX_NEW_TOKENS = 256
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

MODEL_OPTIONS = [
    "sarvamai/OpenHathi-7B-Hi-v0.1-Base",
    "TokenBender/Navarna_v0_1_OpenHermes_Hindi"
]

models = {}
tokenizers = {}

# Custom chat templates
MISTRAL_TEMPLATE = """<s>[INST] {instruction} [/INST]
{response}
</s>
<s>[INST] {instruction} [/INST]
"""

LLAMA_TEMPLATE = """<s>[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.

If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
<</SYS>>

{instruction} [/INST]
{response}
</s>
<s>[INST] {instruction} [/INST]
"""

for model_id in MODEL_OPTIONS:
    tokenizers[model_id] = AutoTokenizer.from_pretrained(model_id)
    models[model_id] = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        load_in_8bit=True,
    )
    models[model_id].eval()
    
    # Set custom chat templates
    if "Navarna" in model_id:
        tokenizers[model_id].chat_template = MISTRAL_TEMPLATE
    elif "OpenHathi" in model_id:
        tokenizers[model_id].chat_template = LLAMA_TEMPLATE

# Initialize Flask app
app = Flask(__name__)

@app.route('/log', methods=['POST'])
def log_results():
    data = request.json
    # Here you can implement any additional processing or storage logic
    print("Logged:", json.dumps(data, indent=2))
    return jsonify({"status": "success"}), 200

@spaces.GPU(duration=90)
def generate(
    model_id: str,
    message: str,
    chat_history: List[Tuple[str, str]],
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = 0.7,
    top_p: float = 0.95,
) -> Iterator[str]:
    model = models[model_id]
    tokenizer = tokenizers[model_id]

    conversation = []
    for user, assistant in chat_history:
        conversation.extend([
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ])
    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        temperature=temperature,
        num_beams=1,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)

def compare_models(
    model1_name: str,
    model2_name: str,
    message: str,
    chat_history1: List[Tuple[str, str]],
    chat_history2: List[Tuple[str, str]],
    max_new_tokens: int,
    temperature: float,
    top_p: float,
) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]], List[Tuple[str, str]], List[Tuple[str, str]]]:
    if model1_name == model2_name:
        error_message = [("System", "Error: Please select two different models.")]
        return error_message, error_message, chat_history1, chat_history2

    output1 = "".join(list(generate(model1_name, message, chat_history1, max_new_tokens, temperature, top_p)))
    output2 = "".join(list(generate(model2_name, message, chat_history2, max_new_tokens, temperature, top_p)))

    chat_history1.append((message, output1))
    chat_history2.append((message, output2))

    log_comparison(model1_name, model2_name, message, output1, output2)

    return chat_history1, chat_history2, chat_history1, chat_history2

def log_comparison(model1_name: str, model2_name: str, question: str, answer1: str, answer2: str, winner: str = None):
    log_data = {
        "question": question,
        "model1": {"name": model1_name, "answer": answer1},
        "model2": {"name": model2_name, "answer": answer2},
        "winner": winner
    }
    
    # Send log data to Flask server
    import requests
    try:
        response = requests.post('http://144.24.151.32:5000/log', json=log_data)
        if response.status_code == 200:
            print("Successfully logged to server")
        else:
            print(f"Failed to log to server. Status code: {response.status_code}")
    except requests.RequestException as e:
        print(f"Error sending log to server: {e}")

def vote_better(model1_name, model2_name, question, answer1, answer2, choice):
    winner = model1_name if choice == "Model 1" else model2_name
    log_comparison(model1_name, model2_name, question, answer1, answer2, winner)
    return f"You voted that {winner} performs better. This has been logged."

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    
    with gr.Row():
        with gr.Column():
            model1_dropdown = gr.Dropdown(choices=MODEL_OPTIONS, label="Model 1", value=MODEL_OPTIONS[0])
            chatbot1 = gr.Chatbot(label="Model 1 Output")
        with gr.Column():
            model2_dropdown = gr.Dropdown(choices=MODEL_OPTIONS, label="Model 2", value=MODEL_OPTIONS[1])
            chatbot2 = gr.Chatbot(label="Model 2 Output")
    
    text_input = gr.Textbox(label="Input Text", lines=3)
    
    with gr.Row():
        max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS)
        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7)
        top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, value=0.95)
    
    compare_btn = gr.Button("Compare Models")
    
    with gr.Row():
        better1_btn = gr.Button("Model 1 is Better")
        better2_btn = gr.Button("Model 2 is Better")

    vote_output = gr.Textbox(label="Voting Result")

    compare_btn.click(
        compare_models,
        inputs=[model1_dropdown, model2_dropdown, text_input, chatbot1, chatbot2, max_new_tokens, temperature, top_p],
        outputs=[chatbot1, chatbot2, chatbot1, chatbot2]
    )
    
    better1_btn.click(
        vote_better,
        inputs=[model1_dropdown, model2_dropdown, text_input, chatbot1, chatbot2, gr.Textbox(value="Model 1", visible=False)],
        outputs=[vote_output]
    )
    
    better2_btn.click(
        vote_better,
        inputs=[model1_dropdown, model2_dropdown, text_input, chatbot1, chatbot2, gr.Textbox(value="Model 2", visible=False)],
        outputs=[vote_output]
    )

if __name__ == "__main__":
    # Start Flask server in a separate thread
    flask_thread = Thread(target=app.run, kwargs={"host": "0.0.0.0", "port": 5000})
    flask_thread.start()
    
    # Start Gradio app
    demo.queue(max_size=10).launch()