#%%
import gradio as gr
from dotenv import load_dotenv
from search import search_bm25, search_exact, prepare_data, merge_results
import os
import json
import datetime
from datasets import load_dataset, Dataset

load_dotenv()

data = prepare_data()

HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "budu_search_data_new")

new = Dataset.from_dict({'Введите запрос':[],
                   'output':[]})
new.push_to_hub("etadevosyan/not_flagged_data",
                token=HF_TOKEN)

def search_handler(query: str):
    results, exact_results = (
        search_bm25(query, data),
        search_exact(query, data)
    )
    
    json_results = merge_results(exact_results, results)
    return {'results': json_results}

def create_ui(query):
    recommendations = []
    results = search_handler(query)
    for result in results['results'][:5]:
        recommendations.append(f"<div style='padding: 10px; border-bottom: 1px solid #ddd;'>{result['name']}</div>")
    log_non_flagged_queries(query,recommendations)
    return gr.HTML(f"<div style='max-height: 400px; overflow-y: auto;'>{''.join(recommendations)}</div>")

def log_non_flagged_queries(query,recommendations):
    dataset = load_dataset('etadevosyan/not_flagged_data')
    dataset.add_item({'Введите запрос':query,
                   'output':recommendations})
    dataset.push_to_hub("etadevosyan/not_flagged_data",
                token=HF_TOKEN)

iface = gr.Interface(
    fn=create_ui,
    inputs=gr.Textbox(label="Введите запрос"),
    outputs=gr.HTML(),  # Use HTML to render custom styled output,
    allow_flagging='manual',
    flagging_callback = hf_writer,
    flagging_options = ['Хорошая рекомендаация',
                        'Плохая рекомендаация'],
    title="Поисковая система BUDU",
)

iface.launch()