from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

llm = init_chat_model("gpt-4o-mini", temperature=0)

multi_query_prompt = ChatPromptTemplate.from_messages([
    ("system", """사용자 질문에 대해 3개의 서로 다른 관점의 검색 쿼리를 생성하세요.
각 쿼리는 원본 질문의 다른 측면을 다뤄야 합니다.
한 줄에 하나씩, 번호 없이 작성하세요."""),
    ("human", "{question}"),
])

multi_query_chain = multi_query_prompt | llm | StrOutputParser()

def multi_query_retrieve(question: str, retriever):
    queries = multi_query_chain.invoke({"question": question})

    # 각 쿼리로 검색 후 결과 합치기 (중복 제거)
    all_docs = []
    seen = set()
    for query in queries.strip().split("\n"):
        query = query.strip()
        if not query:
            continue
        results = retriever.invoke(query)
        for doc in results:
            doc_id = hash(doc.page_content[:200])
            if doc_id not in seen:
                seen.add(doc_id)
                all_docs.append(doc)
    return all_docs

# 사용 예시
results = multi_query_retrieve(
    "RAG에서 검색 품질을 높이는 방법은?",
    retriever,
)

# 관점 다양성을 강조한 프롬프트
diverse_prompt = ChatPromptTemplate.from_messages([
    ("system", """사용자 질문에 대해 3개의 검색 쿼리를 생성하세요.

각 쿼리는 다음 관점 중 하나를 취해야 합니다:
1. 개념/정의 관점: 핵심 개념과 정의를 다루는 쿼리
2. 구현/실전 관점: 실제 구현 방법과 코드를 다루는 쿼리
3. 비교/평가 관점: 다른 기법과의 비교나 장단점을 다루는 쿼리

한 줄에 하나씩, 번호 없이 작성하세요."""),
    ("human", "{question}"),
])

decompose_prompt = ChatPromptTemplate.from_messages([
    ("system", """복잡한 질문을 답변에 필요한 하위 질문들로 분해하세요.
각 하위 질문은 독립적으로 검색 가능해야 합니다.
한 줄에 하나씩, 번호 없이 작성하세요."""),
    ("human", "{question}"),
])

decompose_chain = decompose_prompt | llm | StrOutputParser()

def decompose_and_retrieve(question: str, retriever):
    sub_queries = decompose_chain.invoke({"question": question})

    all_docs = []
    seen = set()
    for query in sub_queries.strip().split("\n"):
        query = query.strip()
        if not query:
            continue
        results = retriever.invoke(query)
        for doc in results:
            doc_id = hash(doc.page_content[:200])
            if doc_id not in seen:
                seen.add(doc_id)
                all_docs.append(doc)
    return all_docs

results = decompose_and_retrieve(
    "Self-RAG와 Corrective RAG의 차이점을 비교하고, Adaptive RAG가 이 둘을 어떻게 통합하는지 설명해줘",
    retriever,
)

stepback_prompt = ChatPromptTemplate.from_messages([
    ("system", """주어진 질문에 대해 한 단계 뒤로 물러서서, 더 넓은 범위의 상위 질문을 생성하세요.
상위 질문은 원래 질문에 답하는 데 필요한 배경 지식을 포함해야 합니다."""),
    ("human", "원본 질문: {question}\n\n상위 질문:"),
])

stepback_chain = stepback_prompt | llm | StrOutputParser()

def stepback_retrieve(question: str, retriever):
    stepback_question = stepback_chain.invoke({"question": question})

    # 원본 질문 + 상위 질문 모두로 검색
    original_results = retriever.invoke(question)
    stepback_results = retriever.invoke(stepback_question)

    # 중복 제거 후 합산
    all_docs = []
    seen = set()
    for doc in original_results + stepback_results:
        doc_id = hash(doc.page_content[:200])
        if doc_id not in seen:
            seen.add(doc_id)
            all_docs.append(doc)
    return all_docs

results = stepback_retrieve(
    "Self-RAG에서 Reflection Token의 ISREL은 어떤 역할을 하나요?",
    retriever,
)

# 예시: Step-back + Multi-Query 조합
def combined_retrieve(question: str, retriever):
    # 1. Step-back으로 상위 질문 생성
    stepback_q = stepback_chain.invoke({"question": question})

    # 2. 원본 + 상위 질문 모두에 Multi-Query 적용
    original_results = multi_query_retrieve(question, retriever)
    stepback_results = multi_query_retrieve(stepback_q, retriever)

    # 3. 중복 제거 후 합산
    all_docs = []
    seen = set()
    for doc in original_results + stepback_results:
        doc_id = hash(doc.page_content[:200])
        if doc_id not in seen:
            seen.add(doc_id)
            all_docs.append(doc)
    return all_docs