import gradio as gr
import requests
import os
import json
import random
from elo import update_elo_ratings  # Custom function for ELO ratings
enable_btn = gr.Button.update(interactive=True)

# Load chatbot URLs and model names from a JSON file
with open('chatbot_urls.json', 'r') as file:
    chatbots = json.load(file)

# Initialize or get user-specific ELO ratings
def get_user_elo_ratings(state):
    return state.get('elo_ratings', {model: 1200 for model in chatbots.keys()})

# Read and write ELO ratings to file (thread-safe)
def read_elo_ratings():
    try:
        with open('elo_ratings.json', 'r') as file:
            return json.load(file)
    except FileNotFoundError:
        return {model: 1200 for model in chatbots.keys()}

def write_elo_ratings(elo_ratings):
    with open('elo_ratings.json', 'w') as file:
        json.dump(elo_ratings, file, indent=4)

# Function to get bot response
def get_bot_response(url, prompt):
    payload = {
        "input": {
            "prompt": prompt,
            "sampling_params": {
                "max_new_tokens": 50,
                "temperature": 0.7,
            }
        }
    }
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "authorization": os.environ.get("RUNPOD_TOKEN")
    }
    response = requests.post(url, json=payload, headers=headers)
    return response.json()['output'][0]['generated_text'].replace(prompt, "")

def chat_with_bots(user_input, state):
    bot_names = list(chatbots.keys())
    random.shuffle(bot_names)
    bot1_url, bot2_url = chatbots[bot_names[0]], chatbots[bot_names[1]]
    
    # Update the state with the names of the last bots
    state.update({'last_bots': [bot_names[0], bot_names[1]]})

    bot1_response = get_bot_response(bot1_url, user_input)
    bot2_response = get_bot_response(bot2_url, user_input)

    return bot1_response, bot2_response

def update_ratings(state, winner_index):
    elo_ratings = get_user_elo_ratings()
    bot_names = list(chatbots.keys())
    winner = state['last_bots'][winner_index]
    loser = state['last_bots'][1 - winner_index]
    
    elo_ratings = update_elo_ratings(elo_ratings, winner, loser)
    write_elo_ratings(elo_ratings)
    return f"Updated ELO ratings:\n{winner}: {elo_ratings[winner]}\n{loser}: {elo_ratings[loser]}"

def vote_up_model(state, chatbot):
    update_message = update_ratings(state, 0)
    chatbot.append(update_message)
    return chatbot

def vote_down_model(state, chatbot):
    update_message = update_ratings(state, 1)
    chatbot.append(update_message)
    return chatbot
def user_ask(state, chatbot1, chatbot2, textbox):
    user_input = textbox
    if len(user_input) > 200:
        user_input = user_input[:200]  # Limit user input to 200 characters

    # Updating state with the current ELO ratings
    state["elo_ratings"] = read_elo_ratings()

    # Chat with bots
    bot1_response, bot2_response = chat_with_bots(user_input, state)

    # Update chat history in state
    if "history" not in state:
        state["history"] = []
    state["history"].extend([
        {"role": "user", "content": user_input},
        {"role": "bot1", "content": bot1_response},
        {"role": "bot2", "content": bot2_response}
    ])

    # Keep only the last 10 messages in history
    state["history"] = state["history"][-10:]

    # Format the conversation in ChatML format
    chatml_prompt = format_chatml_prompt(state)

    return state, chatbot1, chatbot2, textbox

# Gradio interface setup
with gr.Blocks() as demo:
    state = gr.State({})

    with gr.Row():
        with gr.Column():
            chatbot1 = gr.Chatbot(label='Model A').style(height=600)
            upvote_btn_a = gr.Button(value="👍 Upvote A")
        
        with gr.Column():
            chatbot2 = gr.Chatbot(label='Model B').style(height=600)
            upvote_btn_b = gr.Button(value="👍 Upvote B")        

    textbox = gr.Textbox(placeholder="Enter your prompt (up to 200 characters)", max_chars=200)
    submit_btn = gr.Button(value="Send")

    textbox.submit(user_ask, inputs=[state, chatbot1, chatbot2, textbox], outputs=[state, chatbot1, chatbot2, textbox])
    submit_btn.click(user_ask, inputs=[state, chatbot1, chatbot2, textbox], outputs=[state, chatbot1, chatbot2, textbox])
    upvote_btn_a.click(vote_up_model, inputs=[state, chatbot1], outputs=[chatbot1])
    upvote_btn_b.click(vote_down_model, inputs=[state, chatbot2], outputs=[chatbot2])

    demo.launch()