import gradio as gr
import requests
import os
import json
from elo import update_elo_ratings  # Custom function for ELO ratings

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

# Load existing ELO ratings or initialize them
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()}

# Update ELO ratings in a file
def write_elo_ratings(elo_ratings):
    with open('elo_ratings.json', 'w') as file:
        json.dump(elo_ratings, file, indent=4)

elo_ratings = read_elo_ratings()

def get_bot_response(url, prompt):
    payload = {
        "input": {
            "prompt": prompt,
            "sampling_params": {
                "max_new_tokens": 16,
                "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()

def chat_with_bots(user_input):
    bot1_url = chatbots[list(chatbots.keys())[0]]
    bot2_url = chatbots[list(chatbots.keys())[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(bot_index):
    global elo_ratings
    bot_names = list(chatbots.keys())
    winner = bot_names[bot_index]
    loser = bot_names[1 - bot_index]
    
    elo_ratings = update_elo_ratings(elo_ratings, winner, loser)
    write_elo_ratings(elo_ratings)
    
    return f"Bot 1 is {bot_names[0]} with ELO {elo_ratings[bot_names[0]]}\nBot 2 is {bot_names[1]} with ELO {elo_ratings[bot_names[1]]}"

# Gradio Interface
iface = gr.Interface(
    fn=chat_with_bots,
    inputs=[
        gr.inputs.Textbox(label="Your message")
    ],
    outputs=[
        gr.outputs.Textbox(label="Bot 1 Response"),
        gr.outputs.Textbox(label="Bot 2 Response")
    ],
    live=True
)

iface.add_component(gr.inputs.Radio(["Bot 1", "Bot 2"], label="Vote for the Best Response"), "vote")
iface.add_component(gr.outputs.Textbox(label="Voting Result"), "vote_result")

iface.update(value=update_ratings, component_name="vote", outputs=["vote_result"])

iface.launch()