import json
import streamlit as st
from bestrag import BestRAG
import os

# Streamlit app title
col1, col2 = st.columns([1, 5])
with col1:
    st.image("https://github.com/user-attachments/assets/e23d11d5-2d7b-44e2-aa11-59ddcb66bebc", width=140)
with col2:
    st.title("BestRAG - Hybrid Retrieval-Augmented Generation (RAG)")

st.markdown("""
[![GitHub stars](https://img.shields.io/github/stars/samadpls/BestRAG?color=red&label=stars&logoColor=black&style=social)](https://github.com/samadpls/BestRAG)
[![PyPI - Downloads](https://img.shields.io/pypi/dm/bestrag?style=social)](https://pypi.org/project/bestrag/)

> **Note**: Qdrant offers a free tier with 4GB of storage. To generate your API key and endpoint, visit [Qdrant](https://qdrant.tech/).

You can use BestRAG freely by installing it with `pip install bestrag`. For more details, visit the [GitHub repository](https://github.com/samadpls/BestRAG).

Made with ❤️ by [samadpls](https://github.com/samadpls)
""")

# Input fields for BestRAG initialization
url = st.text_input("Qdrant URL", "https://YOUR_QDRANT_URL")
api_key = st.text_input("Qdrant API Key", "YOUR_API_KEY")
collection_name = st.text_input("Collection Name", "YOUR_COLLECTION_NAME")

# Initialize BestRAG only when the user clicks a button
if st.button("Initialize BestRAG"):
    st.session_state['rag'] = BestRAG(url=url, api_key=api_key, collection_name=collection_name)
    st.success("BestRAG initialized successfully!")

# Check if BestRAG is initialized
if 'rag' in st.session_state:
    rag = st.session_state['rag']

    # Tabs for different functionalities
    tab1, tab2 = st.tabs(["Create Embeddings", "Search Embeddings"])

    with tab1:
        st.header("Create Embeddings")
        
        # File uploader for PDF
        pdf_file = st.file_uploader("Upload PDF", type=["pdf"])
        
        if st.button("Create Embeddings"):
            if pdf_file is not None:
                # Save the uploaded PDF to a temporary file
                temp_pdf_path = os.path.join("/tmp", pdf_file.name)
                with open(temp_pdf_path, "wb") as f:
                    f.write(pdf_file.getbuffer())
                
                # Use the uploaded PDF's name
                pdf_name = pdf_file.name
                
                # Store PDF embeddings
                rag.store_pdf_embeddings(temp_pdf_path, pdf_name)
                st.success(f"Embeddings created for {pdf_name}")
            else:
                st.error("Please upload a PDF file.")

    with tab2:
        st.header("Search Embeddings")
        
        # Input fields for search
        query = st.text_input("Search Query", "example query")
        limit = st.number_input("Limit", min_value=1, max_value=20, value=5)
        
        if st.button("Search"):
            # Perform search
            results = rag.search(query, limit)
            
            # Display results
            st.subheader("Search Results")
            for result in results.points:
                st.json({
                    "id": result.id,
                    "score": result.score,
                    "payload": result.payload
                })
else:
    st.warning("Please initialize BestRAG first.")