import os
import gradio as gr
import pyarabic.araby as araby
import pandas as pd
import re
from datasets import load_dataset
from datasets import Features
from datasets import Value
from datasets import Dataset


Secret_token = os.getenv('HF_token')


books = load_dataset('FDSRashid/Hadith_info', data_files='Books.csv', token=Secret_token)['train'].to_pandas()
books.loc[:, 'book_type'] = books['book_type'].apply(lambda x: x.replace('"', '')[2:])
books.loc[:, 'Book_Name'] = books['Book_Name'].apply(lambda x: x.replace('"', ''))
books.loc[:, 'Author'] = books['Author'].apply(lambda x: x.replace('"', ''))
books = books.drop(columns = ['field5',	'field6',	'field7',	'field8'])
css = """
.table-wrap {
    min-height: 300px;
    max-height: 300px;
}

"""

def book_retriever(name):
    if 'ALL' in name:
        return books
    else:
        full_names = name.replace(', ', '|').replace(',', '|')
        return books[(books['Book_Name'].apply(lambda x: araby.strip_diacritics(x)).str.contains(araby.strip_diacritics(name), regex=True)) | (books['Author'].apply(lambda x: araby.strip_diacritics(x)).str.contains(araby.strip_diacritics(name), regex=True)) | (books['Book_ID'].astype(str).isin(full_names.split('|')))]


with gr.Blocks(css=css) as demo:
    text_input = gr.Textbox()
    text_output = gr.DataFrame(wrap=True)
    text_button = gr.Button("Search")
    text_button.click(book_retriever, inputs=text_input, outputs=text_output)

demo.launch()