from dotenv import load_dotenv

load_dotenv()  # .env 에서 OPENAI_API_KEY 로드

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 검색기(retriever) 구성: nvme.txt + ssd.txt → Chroma → as_retriever
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
_docs = TextLoader("../DATA/nvme.txt", encoding="utf-8").load() \
      + TextLoader("../DATA/ssd.txt",  encoding="utf-8").load()
_chunks = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(_docs)
_store = Chroma.from_documents(_chunks, embeddings, collection_name="rag_chain_demo")
retriever = _store.as_retriever(search_kwargs={"k": 3})

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

# (1) 정답
prompt = ChatPromptTemplate.from_messages([
    ("system",
        "당신은 문서 기반 QA 시스템입니다. 아래 문서만 참고해서 답하세요. "
        "문서에 없으면 '모르겠습니다'라고 답하세요.\n\n문서:\n{context}"),
    ("user", "{question}"),
])

def format_docs(docs):
    # (1) 정답
    return "\n\n".join(d.page_content for d in docs)

# (1) 정답
chain = (
    RunnablePassthrough.assign(context=lambda x: format_docs(retriever.invoke(x["question"])))
    | prompt
    | llm
    | StrOutputParser()
)

print(chain.invoke({"question": "NVMe 와 SATA SSD 의 차이?"}))

def format_docs_with_index(docs):
    # (2) 정답
    """청크에 [1], [2] 번호를 붙여서 인용 가능하게"""
    return "\n\n".join(f"[{i}] {d.page_content}" for i, d in enumerate(docs, 1))

def extract_sources(docs):
    # (2) 정답
    """검색된 청크의 출처를 중복 제거해서 리스트로"""
    seen, sources = set(), []
    for d in docs:
        src = d.metadata.get("source", "unknown")
        if src not in seen:
            seen.add(src)
            sources.append(src)
    return sources

def retrieve_and_split(inputs):
    # (2) 정답
    docs = retriever.invoke(inputs["question"])
    return {
        "question": inputs["question"],
        "context":  format_docs_with_index(docs),
        "sources":  extract_sources(docs),
    }

def append_sources(d):
    # (2) 정답
    src_lines = "\n".join(f"  - {s}" for s in d["sources"])
    return f"{d['answer']}\n\n📚 참고 문서:\n{src_lines}"

# (2) 정답
chain = (
    RunnableLambda(retrieve_and_split)
    | RunnablePassthrough.assign(
        answer=(prompt | llm | StrOutputParser())
    )
    | RunnableLambda(append_sources)
)

print(chain.invoke({"question": "NVMe 의 인터페이스가 뭐고 어떤 장점이 있어?"}))

# (3) 정답
# ★ 프롬프트가 '인라인 인용' 을 명시적으로 지시 — 이게 4.2 와의 핵심 차이
prompt = ChatPromptTemplate.from_messages([
    ("system",
        "다음 번호가 매겨진 문서들만 참고해 한국어로 답하세요.\n"
        "각 문장 끝에, 근거가 된 문서 번호를 [1], [2] 처럼 표기하세요.\n"
        "한 문장이 여러 문서를 근거로 하면 [1][2] 처럼 모두 표기합니다.\n"
        "문서에 없으면 '모르겠습니다'.\n\n"
        "문서:\n{context}"),
    ("user", "{question}"),
])

def numbered_sources(docs):
    """청크 번호별 출처 — 중복 제거 안 함 (인라인 [n] 과 1:1 대응)"""
    return [f"[{i}] {d.metadata.get('source', 'unknown')}" for i, d in enumerate(docs, 1)]

def retrieve_and_split(inputs):
    docs = retriever.invoke(inputs["question"])
    return {
        "question": inputs["question"],
        "context":  format_docs_with_index(docs),
        "sources":  numbered_sources(docs),
    }

chain = (
    RunnableLambda(retrieve_and_split)
    | RunnablePassthrough.assign(answer=(prompt | llm | StrOutputParser()))
    | RunnableLambda(append_sources)
)

print("\n(3) 인라인 인용 버전")
print(chain.invoke({"question": "NVMe 의 인터페이스가 뭐야?"}))
