import gradio as gr
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from openai import OpenAI
from translate_utils import translate_ko_to_en

YOUR_OPENAI_API_KEY = "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"

# 임베딩 모델 로드
embedding_model_name = "snunlp/KR-SBERT-V40K-klueNLI-augSTS"
embedding_model = HuggingFaceEmbeddings(
    model_name=embedding_model_name,
    model_kwargs={"device": "cpu"},  # GPU 사용 시 "cuda"로 변경
    encode_kwargs={"normalize_embeddings": True},
)

# vector DB 로드
save_path = "./version-2024-12-22"
vectorstore = FAISS.load_local(save_path, embedding_model, allow_dangerous_deserialization=True)

def chatbot(input_question, eng_trans=True, num_ref=3):
    """챗봇 함수"""
    retriever = vectorstore.as_retriever(search_kwargs={"k": num_ref})

    # 한국어로 질의
    if not eng_trans:
        basic_docs = retriever.invoke(input_question)
    # 영어로 번역 후 질의(영어, 한글 둘 다 검색)
    else:
        eng = translate_ko_to_en(input_question)
        basic_docs = retriever.invoke(input_question)
        eng_docs = retriever.invoke(eng)
        basic_docs = basic_docs + eng_docs

    context = "\n".join([doc.page_content for doc in basic_docs])

    client = OpenAI(api_key=YOUR_OPENAI_API_KEY)

    # GPT-4 or GPT-4o-mini 등 모델 지정
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": f"""당신은 반도체와 차세대반도체학과에 대해서 설명하는 Assistant입니다.
    차세대반도체학과는 한국 서울의 중앙대학교의 창의ICT공과대학에 설립된 학과입니다. 이 학과는 서울대학교, 중앙대학교, 포항공과대학교, 숭실대학교, 강원대학교, 대구대학교, 조선이공대학교가 학점교류를 통해 수강할 수 있는 학과입니다.
    다음 맥락에 맞게 질문에 한글로 답하세요. 대신, 반도체 전문용어는 영어로 해도 됩니다.

    맥락: {context}

    질문: {input_question}
    """
            }
        ]
    )

    return completion.choices[0].message.content

# Gradio Blocks 레이아웃으로 꾸민 버전
with gr.Blocks() as demo:
    # 상단 이미지
    gr.Image(
        value="head.png",  # 이미지 경로 또는 URL
        elem_id="top-image",
        label=None
    )
    
    # 안내 문구
    gr.Markdown(
        """
        # 차세대반도체학과 특화 챗봇
        - 이 챗봇은 한국어 질의에 대해 반도체 관련 정보, 차세대반도체학과 관련 정보를 친절하게 제공합니다.<br>
        - 반도체 전문용어는 일부 영어로 답변될 수 있습니다.<br>
        ### 영어 번역 사용 여부 기능
        - 영어 번역 사용 여부를 키면, 정확도가 상승될 수 있으나, 추론 시간이 길어질 수 있습니다. <br>
        ### 검색 문서 개수 변경 기능
        검색할 문서 개수를 늘리면 정확도가 상승될 수 있으나, 추론 시간이 길어질 수 있습니다. <br>
        """,
        elem_id="description"
    )

    # 메인 UI
    with gr.Group():
        with gr.Row():
            # 입력 파트
            with gr.Column():
                input_question = gr.Textbox(
                    label="질문 입력",
                    placeholder="반도체와 차세대반도체학과에 대해 궁금한 점을 입력하세요."
                )
                eng_trans = gr.Checkbox(
                    label="영어 번역 사용 여부",
                    value=True
                )
                num_ref = gr.Slider(
                    minimum=1, 
                    maximum=5, 
                    value=3, 
                    step=1, 
                    label="검색할 문서 개수"
                )
                submit_btn = gr.Button("질문하기")

            # 출력 파트
            with gr.Column():
                output_answer = gr.Textbox(
                    label="답변",
                    placeholder="답변이 여기에 표시됩니다...",
                    lines=10
                )

        # 버튼과 함수 연결
        submit_btn.click(
            fn=chatbot, 
            inputs=[input_question, eng_trans, num_ref], 
            outputs=output_answer
        )

demo.launch()