import os
os.environ["HF_HOME"] = "/home/user/.cache/huggingface"

from langchain_community.llms import HuggingFacePipeline
import torch


# beomi/gemma-ko-2b 모델 지정
model_id = "beomi/gemma-ko-2b"
llm = HuggingFacePipeline.from_model_id(
    model_id=model_id,
    task="text-generation",
    device=None,  # Explicitly set to None to avoid conflicts with device_map
    model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "auto"}
)

# %%
import os

os.environ['TRANSFORMERS_CACHE'] = '/data/heesu/huggingface_cache'
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
import gradio as gr  # Gradio 임포트 추가

# --- 1. Vector DB 및 임베딩 모델 불러오기 ---
print("Vector DB와 임베딩 모델을 불러오는 중입니다...")
index_path = "proj2_voca"
model_name = "jhgan/ko-sroberta-multitask"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
embeddings = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)
vectorstore = FAISS.load_local(index_path, embeddings, allow_dangerous_deserialization=True)
retriever = vectorstore.as_retriever(search_kwargs={'k': 2})

# --- 2. 언어 모델(LLM) 불러오기 ---
print("언어 모델(beomi/gemma-ko-2b)을 불러오는 중입니다. 시간이 다소 걸릴 수 있습니다...")
model_id = "beomi/gemma-ko-2b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map=None  # 자동 배치 비활성화
).to("cpu")  # GPU 번호 명시

# --- 2. 언어 모델(LLM) 불러오기 ---
# ... (이전 코드 생략) ...
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=-1,
    max_new_tokens=170,  # 최대 생성 토큰 수 감소
    temperature=0.7,
    repetition_penalty=1.2,  # 반복 방지 패널티 추가
    pad_token_id=tokenizer.eos_token_id
)
# ... (이후 코드 생략) ...
llm = HuggingFacePipeline(pipeline=pipe)

# --- 3. 프롬프트 템플릿 정의 ---
template = """
당신은 사용자의 꿈을 명확하게 해석해주는 꿈 해몽 전문가입니다. 아래 '검색된 꿈 해몽 정보'를 바탕으로, 사용자의 질문에 대해 가장 관련성이 높은 내용을 한 문장으로 요약해서 답변해주세요


### 검색된 꿈 해몽 정보:
{context}



### 사용자의 질문:
{question}



### 전문가의 답변:
"""
prompt = PromptTemplate.from_template(template)

# --- 4. RAG 체인(Chain) 구성 ---
rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
)


# --- 응답 청소 함수 추가 (원본에서 가져옴) ---
def clean_response(response):
    # "###"로 분할하고 첫 번째 부분만 반환 (불필요한 후속 내용 제거)
    cleaned = response.split("###")[0].strip()
    return cleaned


# --- 5. 꿈 해몽 함수 정의 (Gradio에서 호출될 함수) ---
def interpret_dream(query):
    docs_with_scores = vectorstore.similarity_search_with_score(query, k=2)
    score_threshold = 1.2
    filtered_docs = [doc for doc, score in docs_with_scores if score < score_threshold]

    if not filtered_docs:
        return "죄송합니다. 현재 데이터베이스에 해당 꿈에 대한 정보가 부족하여 해몽을 제공하기 어렵습니다."
    else:
        context_texts = [f"- {d.page_content}: {d.metadata['meaning']}" for d in filtered_docs]

        local_rag_chain = (
                prompt
                | llm
                | StrOutputParser()
        )

        response = local_rag_chain.invoke({
            "context": "\n".join(context_texts),
            "question": query
        })

        try:
            final_answer = response.split("### 전문가의 답변:")[1].strip()
        except IndexError:
            final_answer = response.strip()

        # 청소 함수 적용 (원본에서 가져옴)
        final_answer = clean_response(final_answer)

        return final_answer


# --- 6. Gradio 인터페이스 구성 ---
with gr.Blocks() as demo:
    gr.Markdown("# ✨ 꿈 해몽 서비스")
    gr.Markdown("당신의 꿈 내용을 입력하세요. AI가 명확하게 해석해 드립니다!")

    with gr.Row():
        input_text = gr.Textbox(label="당신의 꿈은 무엇이었나요?", placeholder="예: 파란 불이 나는 꿈")
        output_text = gr.Textbox(label="꿈 해몽 결과", interactive=False)

    submit_button = gr.Button("해몽하기")
    submit_button.click(interpret_dream, inputs=input_text, outputs=output_text)

# 앱 실행
if __name__ == '__main__':
    demo.launch(share=True)  # share=True로 공공 URL 생성 (선택사항)