import gradio as gr
import requests
import json
import os
from datasets import load_dataset

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

db_password = os.environ["db_password"]

uri = f"mongodb+srv://tuna:{db_password}@tuna.ixw4ff2.mongodb.net/?appName=tuna"

# Create a new client and connect to the server
client = MongoClient(uri, server_api=ServerApi('1'))

database = client['valid_image_caption']
collection = database["log_valid"]

ds = load_dataset("anhdt-dsai-02/test_image_dataset_1_2_3_4", token = os.environ['token_huggingface'] )

def get_similar_captions(caption):

    url = "https://anhdt-dsai-02-caption-retrieval.hf.space/retrieval"
    params = {"caption": caption}

    try:
        response = requests.post(url, params=params, headers={"accept": "application/json"})
        response.raise_for_status()  # Raises an HTTPError for bad responses (4xx and 5xx)
        data = response.json()  # Assuming response is JSON

    except requests.exceptions.RequestException as e:
        print(f"Request failed: {e}")
    except ValueError:
        print("Failed to parse JSON response")

    return data



index = 0

def create_question():
    global index
    print(index)
    caption = ds['train'][index]["caption"] # delete [0]
    caption_lst = get_similar_captions(caption)
    #id_lst = [pair['id'] for pair in caption_lst]

    id_lst = [int(pair['id']/10) for pair in caption_lst]

    images = [ds["train"][id]['image'] for id in id_lst]

    index += 1

    return caption, images, id_lst

def on_select(evt: gr.SelectData):
    idx = evt.index
    return idx

# Function to handle user selection
def process_answer(user_id, caption, selected_image, ids):

    collection.insert_one({ "used_id" : user_id, "caption": caption, "ids":  ids, "choice": selected_image})
    send_score(user_id, 0.1)

    return create_question()

def send_score(user_id, score = 0.1):
    max_retries = 10
    while max_retries > 0:
        url = os.environ['api_url'] + "grade"

        payload = {
            "token": user_id,
            "comment": "Good job!",
            "grade": score,
            "submitted_at": "2021-01-01 00:00:00",
            "graded_at": "2021-01-01 00:00:00"
        }
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json",
            "X-Public-Api-Key": os.environ['ADMIN']
        }

        response = requests.post(url, json=payload, headers=headers)
        if response.status_code == 200:
            return True
        print(response)
        max_retries -= 1
    return False

def authenticate(user_id):

    url = os.environ['api_url'] + "authenticate"
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
        "X-Public-Api-Key": os.environ['ADMIN']
        }
    payload = { "token": user_id }
    response = requests.post(url, json=payload, headers=headers)

    return response.status_code == 200



def login(username):
    #state[0] = username
    #package[0] = get_next_package(user_id=username)
    """
    #temp
    gr.Info("Login successfully. Welcome!")
    return f"Welcome, {username}!", gr.update(visible=False), gr.update(visible=True)
    #temp
    """

    # Authenticate user
    if authenticate(username):
        #user_sessions[username] = True
        gr.Info("Login successfully. Welcome!")
        return f"Welcome, {username}!", gr.update(visible=False), gr.update(visible=True), *create_question()
    else:
        raise gr.Error("Token ID is invalid! Try again!")
        return "Invalid Token ID", gr.update(visible=True), gr.update(visible=False)

# Gradio UI
with gr.Blocks() as demo:
    with gr.Column(visible=True) as login_section:
        username_input = gr.Textbox(placeholder="Enter your token", label="Token ID", type="password")
        login_button = gr.Button("Login")
        login_output = gr.Textbox(label="Login Status", interactive=False)


    # Upload section  (initially hidden)
    with gr.Column(visible=False) as upload_section:
        caption = gr.State()
        images = gr.State()
        id_lst = gr.State()
        # Store the selected image path
        selected_image = gr.State()


        #caption.value, images.value, id_lst.value = create_question()

        gr.Markdown("**Đâu là hình ảnh phù hợp với caption sau:**")
        markdown = gr.Markdown(caption.value)



        # Display images in a Gallery
        gallery = gr.Gallery(value=images.value, label="Ấn vào hình ảnh để chọn sau đó Submit", columns = 4)

        # Use the Gallery's select event to update the selected image
        gallery.select(fn=on_select, outputs=selected_image)

        # Submit button
        submit_button = gr.Button("Submit")

        # Output message
        output_message = gr.Textbox(label="Result")

    # Action in Login Section
    username_input.submit(
        login, inputs=[username_input], outputs=[login_output, login_section, upload_section,  markdown, gallery, id_lst] #, translation_section, en_input, vi_input]
    )
    login_button.click(
        login, inputs=[username_input], outputs=[login_output, login_section, upload_section,  markdown, gallery, id_lst] #, translation_section, en_input, vi_input]
    )

    # Link the submit button to the processing function

    submit_button.click(fn=process_answer, inputs=[username_input, caption, selected_image, id_lst], outputs=[markdown, gallery, id_lst])
# Launch the app
demo.launch(debug = True)