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

import sqlite3

from pymongo.mongo_client import MongoClient
from pymongo.server_api import ServerApi

password=os.environ.get("MONGODB")

def init_database():
    uri = f"mongodb+srv://new-user:{password}@cluster0.xb2urf6.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0"
    client = MongoClient(uri, ssl_cert_reqs=ssl.CERT_NONE)
    db = client["elo_ratings"]
    collection = db["elo_ratings"]
    return collection

def get_user_elo_ratings(collection):
    rows = list(collection.find())
    if rows:
        return {row['bot_name']: {'elo_rating': row['elo_rating'], 'games_played': row['games_played']} for row in rows}
    else:
        return {"default": {'elo_rating': 1200, 'games_played': 0}}

def update_elo_rating(collection, updated_ratings, winner, loser):
    collection.update_one({"bot_name": winner}, {"$set": {"elo_rating": updated_ratings[winner]['elo_rating'], "games_played": updated_ratings[winner]['games_played']}}, upsert=True)
    collection.update_one({"bot_name": loser}, {"$set": {"elo_rating": updated_ratings[loser]['elo_rating'], "games_played": updated_ratings[loser]['games_played']}}, upsert=True)

# Load chatbot URLs and model names from a JSON file
# Load chatbot model adapter names from a text file
with open('chatbots.txt', 'r') as file:
    chatbots = file.read().splitlines()
def clear_chat(state):
    # Reset state including the chatbot order
    state = {} if state is not None else state
    # Shuffle and reinitialize chatbots in the state
    bot_names = list(chatbots.keys())
    random.shuffle(bot_names)
    state['last_bots'] = [bot_names[0], bot_names[1]]
    # Reset other components
    return state, None, None, gr.Button.update(interactive=False), gr.Button.update(interactive=False), gr.Textbox.update(interactive=True), gr.Button.update(interactive=True)

global_elo_ratings=None

from datasets import load_dataset,DatasetDict,Dataset
import requests
import os


# Function to get bot response
def format_alpaca_prompt(state):
    alpaca_prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
    alpaca_prompt2 = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
    for message in state["history"][0]:
        j=""
        if message['role']=='user':
            j="### Instruction:\n"
        else:
            j="### Response:\n"
        alpaca_prompt += j+ message['content']+"\n\n"
    for message in state["history"][1]:
        j=""
        if message['role']=='user':
            j="### Instruction:\n"
        else:
            j="### Response:\n"
        alpaca_prompt2 += j+ message['content']+"\n\n"
    return [alpaca_prompt+"### Response:\n",alpaca_prompt2+"### Response:\n"]
import aiohttp
import asyncio
from tenacity import retry, stop_after_attempt, wait_exponential

async def get_bot_response(adapter_id, prompt, state, bot_index):
    alpaca_prompt = format_alpaca_prompt(state)
    print(alpaca_prompt)
    payload = {
        "inputs": alpaca_prompt[bot_index],
        "parameters": {
            "adapter_id": adapter_id,
            "adapter_source": "hub",
            "temperature": 1,
            "max_new_tokens": 100
        }
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ.get('PREDIBASE_TOKEN')}"
    }
    async with aiohttp.ClientSession() as session:
        try:
            async with session.post("https://serving.app.predibase.com/79957f/deployments/v2/llms/mistral-7b-instruct/generate",
                                    json=payload, headers=headers, timeout=30) as response:
                if response.status == 200:
                    response_data = await response.json()
                    response_text = response_data.get('generated_text', '')
                else:
                    print(response.text)
                    response_text = "Sorry, I couldn't generate a response."
        except (aiohttp.ClientError, asyncio.TimeoutError):
            response_text = "Sorry, I couldn't generate a response."
    return response_text.split('### Instruction')[0]
async def chat_with_bots(user_input, state):
    # Use existing bot order from state if available, otherwise shuffle and initialize
    if 'last_bots' not in state or not state['last_bots']:
        random.shuffle(chatbots)
        state['last_bots'] = [chatbots[0], chatbots[1]]
    
    bot1_adapter, bot2_adapter = state['last_bots'][0], state['last_bots'][1]

    bot1_response, bot2_response = await asyncio.gather(
        get_bot_response(bot1_adapter, user_input, state, 0),
        get_bot_response(bot2_adapter, user_input, state, 1)
    )

    return bot1_response, bot2_response
def update_ratings(state, winner_index, collection):
    elo_ratings = get_user_elo_ratings(collection)
    winner = state['last_bots'][winner_index]
    loser = state['last_bots'][1 - winner_index]
    
    elo_ratings = update_elo_ratings(elo_ratings, winner, loser)
    update_elo_rating(collection, elo_ratings, winner, loser)
    return [('Winner: ', winner.replace('rwitz/','').replace('-lora','')), ('Loser: ', loser.replace('rwitz/','').replace('-lora',''))]
    
def vote_up_model(state, chatbot, chatbot2):
    collection = init_database()
    update_message = update_ratings(state, 0, collection)
    chatbot.append(update_message[0])
    chatbot2.append(update_message[1])
    return chatbot, chatbot2, gr.Button.update(interactive=False), gr.Button.update(interactive=False), gr.Textbox.update(interactive=False), gr.Button.update(interactive=False)

def vote_down_model(state, chatbot, chatbot2):
    collection = init_database()
    update_message = update_ratings(state, 1, collection)
    chatbot2.append(update_message[0])
    chatbot.append(update_message[1])
    return chatbot, chatbot2, gr.Button.update(interactive=False), gr.Button.update(interactive=False), gr.Textbox.update(interactive=False), gr.Button.update(interactive=False)

async def user_ask(state, chatbot1, chatbot2, textbox):
    global enable_btn
    user_input = textbox
    if len(user_input) > 200:
        user_input = user_input[:200]  # Limit user input to 200 characters

    collection = init_database()  # Initialize the collection object

    # Updating state with the current ELO ratings
    state["elo_ratings"] = get_user_elo_ratings(collection)
    if "history" not in state:
        state.update({'history': [[],[]]})
    state["history"][0].extend([
        {"role": "user", "content": user_input}])
    state["history"][1].extend([
        {"role": "user", "content": user_input}])
    # Chat with bots
    bot1_response, bot2_response = await chat_with_bots(user_input, state)
    state["history"][0].extend([
        {"role": "bot1", "content": bot1_response},
    ])
    state["history"][1].extend([
        {"role": "bot2", "content": bot2_response},
    ])
    
    chatbot1.append((user_input,bot1_response))
    chatbot2.append((user_input,bot2_response))

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

    # Format the conversation in ChatML format

    return state, chatbot1, chatbot2, gr.update(value=''),enable_btn,enable_btn
import pandas as pd

# Function to generate leaderboard data

def generate_leaderboard(collection):
    rows = list(collection.find())
    leaderboard_data = pd.DataFrame(rows, columns=['bot_name', 'elo_rating', 'games_played'])
    leaderboard_data.columns = ['Chatbot', 'ELO Score', 'Games Played']
    return leaderboard_data
    
def refresh_leaderboard():
    collection = init_database()
    return generate_leaderboard(collection)

# ...

# Gradio interface setup
# Gradio interface setup
with gr.Blocks() as demo:
    state = gr.State({})
    with gr.Tab("Chatbot Arena"):
        with gr.Row():
            with gr.Column():
                chatbot1 = gr.Chatbot(label='Model A').style(height=600)
                upvote_btn_a = gr.Button(value="👍 Upvote A", interactive=False)
            
            with gr.Column():
                chatbot2 = gr.Chatbot(label='Model B').style(height=600)
                upvote_btn_b = gr.Button(value="👍 Upvote B", interactive=False)        
    
        textbox = gr.Textbox(placeholder="Enter your prompt (up to 200 characters)")
        with gr.Row():
            submit_btn = gr.Button(value="Send")
            reset_btn = gr.Button(value="Reset")
        reset_btn.click(clear_chat, inputs=[state], outputs=[state, chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
        textbox.submit(user_ask, inputs=[state, chatbot1, chatbot2, textbox], outputs=[state, chatbot1, chatbot2, textbox, upvote_btn_a, upvote_btn_b], queue=True)
        submit_btn.click(user_ask, inputs=[state, chatbot1, chatbot2, textbox], outputs=[state, chatbot1, chatbot2, textbox, upvote_btn_a, upvote_btn_b], queue=True)
        
        collection = init_database()
        upvote_btn_a.click(vote_up_model, inputs=[state, chatbot1, chatbot2], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
        upvote_btn_b.click(vote_down_model, inputs=[state, chatbot1, chatbot2], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
    with gr.Tab("Leaderboard"):
        try:
            leaderboard = gr.Dataframe(refresh_leaderboard(collection))
        except:
            leaderboard = gr.Dataframe()
        refresh_btn = gr.Button("Refresh Leaderboard")
        refresh_btn.click(refresh_leaderboard, outputs=[leaderboard])

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch(share=False)