import gradio as gr
from sentence_transformers import SentenceTransformer
import datasets
import time
import faiss

# Impor tema custom dari themes.py
from themes_2 import ClassicMinimalTheme


# Inisialisasi dataset dan model hanya sekali di Gradio
def initialize():
    dataset = datasets.load_dataset('A-Roucher/english_historical_quotes', download_mode="force_redownload")['train']
    all_authors = list(set(dataset['author']))
    model_name = "BAAI/bge-small-en-v1.5"
    encoder = SentenceTransformer(model_name)
    index = faiss.read_index('index_alone.faiss')
    
    return dataset, encoder, index, all_authors

# Fungsi untuk melakukan pencarian kutipan berdasarkan query
def search(query, dataset, encoder, index, progress=gr.Progress()):
    start = time.time()

    if len(query.strip()) == 0:
        return "Silakan masukkan ide atau kata kunci."  # Tidak ada pencarian jika query kosong

    # Memulai progres dengan 0%
    progress(0, desc="Memulai proses...")

    # Encode query menjadi embedding
    progress(0.2, desc="Proses encoding...")
    query_embedding = encoder.encode([query])
    time.sleep(1)  # Simulasi proses lambat

    # Cari kutipan yang paling mirip menggunakan faiss
    progress(0.5, desc="Menemukan kutipan yang mirip...")
    _, samples = index.search(query_embedding, k=10)
    time.sleep(1)

    quotes = dataset.select(samples[0])

    result = "\n\n"
    for i in progress.tqdm(range(len(quotes)), desc="Menyusun hasil..."):
        time.sleep(0.25)  # Simulasi penundaan saat menyusun hasil
        # Menggunakan HTML untuk mengatur ukuran font pada quote dan tebal pada author
        result += f"## ⭐ {quotes['author'][i]}\n> {quotes['quote'][i]}\n----\n"
        # result += f"<b>{quotes['author'][i]}</b><br><span style='font-size: 20px;'>{quotes['quote'][i]}</span></hr>"

    delay = "%.3f" % (time.time() - start)
    
    # Proses selesai 100%
    progress(1.0, desc="Selesai!")

    return f"_Waktu komputasi: **{delay} detik**_{result}"

# Fungsi untuk memulai pencarian dengan progress tracking
def run_search(query):
    dataset, encoder, index, _ = initialize()  # Ambil state inisialisasi (dataset, encoder, index)
    return search(query, dataset, encoder, index)  # Mengembalikan hasil akhir sebagai string

# CSS Kustom untuk mempercantik tampilan
css = """
#header {
    text-align: center;
    color: #FFFFFF;
    background: linear-gradient(90deg, #FF5733, #C70039);
    padding: 20px;
    border-radius: 10px;
}

#input-section, #output-section {
    margin: 10px auto;
    max-width: 700px;
    background-color: #F0F0F0;
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.1);
}

#footer {
    text-align: center;
    margin-top: 30px;
    color: #888;
}

#submit-button {
    background-color: #FF5733;
    color: #FFF;
    font-weight: bold;
    padding: 10px;
    border-radius: 10px;
    cursor: pointer;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
}

body {
    background-color: #F8F9F9;
}

.note {
    background-color: #ddffdd;
    border: 1px solid #ddd;
    padding: 10px;
    margin: 10px 0;
    border-radius: 5px;
    font-family: Arial, sans-serif;
}
"""

# -----------------
# Antarmuka Gradio
# -----------------
with gr.Blocks(css=css, theme=ClassicMinimalTheme()) as Apps:

    # Tambahkan banner
    gr.HTML("""
        <div style='text-align: center;'>
            <img src='https://i.ibb.co.com/x6XxZC9/banner02.jpg' alt='Banner' style='width: 100%; height: auto;'/>
        </div>
    """)

    # Ganti judul
    gr.Markdown("<h1 style='text-align: center;'>Siapa yang memiliki ide serupa dengan Anda?</br>📚 Kami bantu Anda mencarinya dengan AI.. </h1>")
 
    
    # Input Section
    with gr.Row(elem_id="input-section"):
        text_input = gr.Textbox(label="Ketik ide Anda di sini:", placeholder="Money .", lines=2)
        submit_button = gr.Button("🔍 Cari kutipan!", elem_id="submit-button")
    
    # Output Section
    with gr.Row(elem_id="output-section"):
        output = gr.Markdown()

    # Menghubungkan tombol pencarian dengan fungsi search
    submit_button.click(run_search, inputs=[text_input], outputs=[output])

    # Tambahkan footer di bagian bawah
    gr.HTML("""
    <footer id="footer">
        Transfer Energi Semesta Digital © 2024 __drat. | 🇮🇩 Untuk Indonesia Jaya!
    </footer>
    """)


# Meluncurkan aplikasi
if __name__ == "__main__":
    Apps.queue(api_open=False).launch(show_api=False)