import weaviate # vector DB
from openai import OpenAI # LLM
import PyPDF2 # pdf -> text
import numpy as np
from transformers import AutoModel, AutoTokenizer # Model, Tokenzier Load
import gradio as gr # front-end(ui & ux)
from sentence_transformers import SentenceTransformer # embedding

# Weaviate에 사용할 클래스 스키마 정의
def create_schema():
    class_obj = {
        "class": "PdfSentence",  # 클래스명
        "properties": [
            {
                "name": "sentence",
                "dataType": ["text"]
            },
            {
                "name": "embedding",
                "dataType": ["number[]"]  # 벡터 타입
            }
        ]
    }
    # 스키마 생성
    db_client.schema.create_class(class_obj)

# 스키마 확인 및 생성 / 기존에 선언한 스키마가 있는 경우 -> pass(넘어간다.)
def ensure_schema():
    schema = db_client.schema.get()
    classes = [cls["class"] for cls in schema["classes"]]
    print(classes)
    if "PdfSentence" not in classes:
        create_schema()


# 스키마가 존재하지 않을 경우 생성
# if not client.schema.contains({"class": "PdfSentence"}):
#     create_schema()

# PDF 텍스트 추출 함수
def extract_text_from_pdf(pdf):
    pdf_reader = PyPDF2.PdfReader(pdf)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

# 텍스트 임베딩 생성 함수
def create_embeddings(text, model):
    result = model.encode(text)
    return result.astype(np.float64).tolist()


# Weaviate에 데이터 저장
def store_vectors_in_weaviate(sentences, embed_model):
    with db_client.batch as batch:
        for sentence in sentences:
            try:
                embedding = create_embeddings(sentence, embed_model)
                # print("embedding", embedding)

                # Weaviate에 문장과 벡터 저장
                data_object = {
                    "sentence": sentence,
                    "embedding": embedding
                }
                batch.add_data_object(data_object, "PdfSentence")
                print("success")
            except Exception as e:
                print(e)

# 질문에 가장 유사한 문장 찾기
def find_similar_sentence_in_weaviate(question_embedding):
    near_vector = {
        "vector": question_embedding
    }

    result = db_client.query.get("PdfSentence", ["sentence", "embedding"]) \
        .with_near_vector(near_vector) \
        .do()
    #         .with_limit(1) \
    
    return result # 가장 가까이 있는 문장이 반환
    # return result['data']['Get']['PdfSentence'][0]['sentence']

def generate_answer(pdf, question):
    global embed_model
    # PDF 텍스트 추출
    text = extract_text_from_pdf(pdf)
    
    # 텍스트를 문장별로 나누기
    sentences = text.split('. ')
    
    # 문장들을 Weaviate에 저장
    store_vectors_in_weaviate(sentences, embed_model)
    
    # 질문에 대한 임베딩 생성
    question_embedding = create_embeddings(question, embed_model)
    
    # Weaviate에서 가장 유사한 문장 찾기
    most_similar_sentence = find_similar_sentence_in_weaviate(question_embedding)
    print("debug03")
    print(most_similar_sentence)

    # OpenAI API로 응답 생성
    ai_client = OpenAI(api_key='sk-kFppolOqiih2qTFViLZGyuyNx_iw10bcXPTepEUKEKT3BlbkFJbe-U6M4dkZG_IOUJjPjmdyv0OXI2RTl1QZwkCfmV4A')
    response = ai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"You are a helpful assistant. Answer based on context: {most_similar_sentence}"},
            {"role": "user", "content": question}
        ]
    )
    result = response.choices[0].message.content

    return result

# Gradio 인터페이스 구성
def interface(pdf, question):
    return generate_answer(pdf, question)

if __name__=="__main__":
    # 임베딩 모델 로드
    embed_model = SentenceTransformer('xlm-r-100langs-bert-base-nli-stsb-mean-tokens')

    # Weaviate 클라이언트 설정
    db_client = weaviate.Client(
        url="https://qms3ceu4svyrpb8xf36w.c0.asia-southeast1.gcp.weaviate.cloud",  # 클러스터 URL
        auth_client_secret=weaviate.AuthApiKey(api_key="lLiZbr2JMMbiV1PctLfrQeM32s2pC7KSJIms"),  # API 키 설정
        # timeout_config=(5, 150)  # 타임아웃 설정 (선택 사항)
    )

    ensure_schema()

    # Gradio UI 생성
    with gr.Blocks() as demo:
        pdf_input = gr.File(label="Upload PDF", type="filepath")
        question_input = gr.Textbox(label="Ask a question", placeholder="What do you want to know?")
        output = gr.Textbox(label="Answer")

        submit_btn = gr.Button("Submit")
        submit_btn.click(fn=interface, inputs=[pdf_input, question_input], outputs=output)

    # 앱 실행
    demo.launch()