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

# Initialize persistent client
CHROMA_PATH = "./vectordb-chroma/gformdb"
os.makedirs(CHROMA_PATH, exist_ok=True)
default_user_aspek = "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"

def process_files(csv_file, openai_key, user_aspek=default_user_aspek, pdf_file=False, model_name="gpt-4o-mini"):
    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 
        os.environ["OPENAI_API_KEY"] = openai_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: [0-100]
        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:
        {Aspek}


        """
        
        PROMPT = PromptTemplate(
            input_variables=["QA", "context", "Aspek"],
            template=TEMPLATE
        )
        
        # Initialize ChatGPT
        chat = ChatOpenAI(model=model_name)
        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)}"

# 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"]
                    )

            user_aspek = gr.Textbox(
                label="Aspek",
                placeholder="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",
                show_copy_button=True
            )

            model_name = gr.Textbox(
                label="openAI model",
                placeholder="input your openAI model"
            )
            
            api_key = gr.Textbox(
                label="OpenAI API Key",
                placeholder="Enter your OpenAI API key",
                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)