import gradio as gr
import pandas as pd
from langchain.memory import VectorStoreRetrieverMemory
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.prompts import PromptTemplate
import chromadb
import os
from uuid import uuid4
from langchain_google_genai import ChatGoogleGenerativeAI

# Initialize persistent client
CHROMA_PATH = "./vectordb-chroma/gformdb"
os.makedirs(CHROMA_PATH, exist_ok=True)

def process_files(csv_file, gemini_key, user_aspek, pdf_file, model):
    try:
        context = ""

        def context_search(query="a",source_name=False, k=5):

            if source_name==True:
                results = vectorstore.similarity_search(
                    query=query,
                    k=k,
                    filter={"source": source_name}
                )
            else:
                results = vectorstore.similarity_search(
                    query=query,
                    k=k,
                )
            for res in results:
                page_content = res.page_content
                metadata = res.metadata
            return page_content, metadata

        def format_qa_efficient(data_path):
            data = pd.read_csv(data_path)
            data.head()
            qa_efficient = ""
            indexs = 0
            Question = ""
            for column in data.columns[3:]:  # Lewati kolom Timestamp, Email Address, dan Nama
                qa_efficient += f"Nomor-{indexs+1}.{column}:\n"
                indexs += 1
                Question += column
                for index, row in data.iterrows():
                    name = row['Nama']
                    email = row["Email_Address"]
                    answer = row[column]
                    qa_efficient += f"- {name}|{email}: {answer}\n"
                
                qa_efficient += "\n"  # Baris kosong antara setiap pertanyaan

            return qa_efficient, Question

        # Set OpenAI API
        if gemini_key == "":
            gemini_key = os.getenv('GEMINI_API_KEY')
            
        
        # Process CSV (Gradio provides the file path directly)

        QA, Question = format_qa_efficient(csv_file)

        if pdf_file:
            # Initialize Chroma client
            persistent_client = chromadb.PersistentClient(path=CHROMA_PATH)
            collection = persistent_client.get_or_create_collection("RAG")
            
            vectorstore = Chroma(
                client=persistent_client,
                collection_name=collection.name,
                embedding_function=OpenAIEmbeddings()
            )
            
            res, metadata = context_search(pdf_file,k=1)

            if metadata == pdf_file.name:
                # Process PDF (Gradio provides the file path directly)
                pages = PyPDFLoader(pdf_file.name).load_and_split()
                uuids = [str(uuid4()) for _ in range(len(pages))]
                vectorstore.add_documents(documents=pages, ids=uuids)
            context, metadata = context_search(Question, pdf_file.name)

        elif pdf_file == False:
            context = "Tidak ada context tambahan yang diberikan, tolong gunakan pengetahuan anda untuk menjawab pertanyaan"
        
        # Evaluation template
        TEMPLATE = """
        Kamu adalah AI evaluator pendidikan. Nilai jawaban siswa berikut menggunakan panduan ini:

        Buat output dengan format(format ini untuk satu murid, jadi selesaikan dulu 1 murid untuk semua nomor, baru ke murid berikutnya):
        HASIL EVALUASI
        =============
        Nama dan Email Murid: 
        [Nomor Soal].
        Nilai: [Sesuaikan dengan Aspek Penilaian yang diberikan user]
        Alasan: [alasan singkat penilaian]
        Saran Perbaikan: [saran]

        [buat seperti di atas untuk setiap jawaban]

        ================
        Rata-rata Nilai: [nilai]
        
        Rekomendasi Umum:
        [rekomendasi]
        

        MATERI REFERENSI:
        {context}   

        SOAL & JAWABAN:
        {QA}

        Aspek Penilaian yang diberikan user:
        {Aspek}


        """
        
        PROMPT = PromptTemplate(
            input_variables=["QA", "context", "Aspek"],
            template=TEMPLATE
        )
        
        # Initialize ChatGPT
        chat = ChatGoogleGenerativeAI(model=model, api_key=gemini_key)
        chain = PROMPT | chat
        
        # Generate evaluation
        Aspek_penilaian = user_aspek
        
        response = chain.invoke({
            "Aspek": Aspek_penilaian,
            "QA": QA,
            "context": context
        })
        
        return response.content
        
    except Exception as e:
        return f"Error occurred: {str(e)}\nType: {type(e)}"

def load_demo_files(checked):
    if checked:
        return gr.update(value="./demo.csv"), gr.update(value="./demo.pdf")
    else:
        return gr.update(value=None), gr.update(value=None)
        
# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Quiz Evaluator", theme=gr.themes.Soft()) as app:
        with gr.Column(scale=1):
            gr.Markdown(
                """
                # 📝 Quiz Response Evaluator
                Upload your quiz responses (CSV) and reference material (PDF) to get AI-powered evaluation.
                """
            )
            
            with gr.Row():
                with gr.Column(scale=1):
                    csv_input = gr.File(
                        label="Quiz Responses (CSV)",
                        file_types=[".csv"]
                    )
                with gr.Column(scale=1):
                    pdf_input = gr.File(
                        label="Reference Material (PDF)",
                        file_types=[".pdf"]
                    )
                    
            demo_checkbox = gr.Checkbox(
                label="Use Demo Files",
                value=False
            )

            demo_checkbox.change(
                fn=load_demo_files,
                inputs=[demo_checkbox],
                outputs=[csv_input, pdf_input]
            )

            user_aspek = gr.Textbox(
                label="Aspek",
                placeholder="isi dengan bagaimana cara AI akan menilai jawaban",
                show_copy_button=True,
                value = "Jika Jawaban salah berikan nilai 0, jika jawaban benar namun tidak tepat berikan nilai 50, jika jawaban benar dan lengkap serta penjelasan baik, beri nilai 100"
            )

            model_name = gr.Textbox(
                label="Gemini Model",
                placeholder="input your Gemini model",
                value="gemini-1.5-pro"
            )
            
            api_key = gr.Textbox(
                label="Gemini API Key",
                placeholder="Enter your Gemini API key, default is my API key, use it wisely XD",
                type="password"
            )
            
            submit_btn = gr.Button(
                "Evaluate Responses", 
                variant="primary",
                size="lg"
            )
            
            # Using Textbox instead of Markdown for better formatting
            output = gr.Textbox(
                label="Evaluation Results",
                lines=20,
                max_lines=30,
                show_copy_button=True
            )

        
        submit_btn.click(
            fn=process_files,
            inputs=[csv_input,api_key,user_aspek,pdf_input,model_name],
            outputs=output
        )
        
    return app

# Launch the interface
if __name__ == "__main__":
    app = create_interface()
    app.launch(share=True)