from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.schema import Document
from pathlib import Path

def load_documents(file_path: str) -> list[Document]:
    """파일 경로에서 문서를 로드합니다."""
    path = Path(file_path)

    if path.suffix.lower() == ".pdf":
        loader = PyPDFLoader(str(path))
    elif path.suffix.lower() in [".txt", ".md"]:
        loader = TextLoader(str(path), encoding="utf-8")
    else:
        raise ValueError(f"지원하지 않는 파일 형식: {path.suffix}")

    documents = loader.load()
    print(f"로드 완료: {len(documents)}개 문서 (출처: {path.name})")
    return documents

def load_multiple_documents(file_paths: list[str]) -> list[Document]:
    """여러 파일에서 문서를 로드합니다."""
    all_docs = []
    for path in file_paths:
        try:
            docs = load_documents(path)
            all_docs.extend(docs)
        except Exception as e:
            print(f"로드 실패 ({path}): {e}")
    print(f"\n총 {len(all_docs)}개 문서 로드 완료")
    return all_docs

# 테스트용 샘플 문서 생성
sample_text = """
Transformer 아키텍처

Transformer는 2017년 Google의 "Attention Is All You Need" 논문에서 제안된 아키텍처입니다.
기존 RNN 기반 Seq2Seq 모델의 한계를 극복하기 위해 Self-Attention 메커니즘을 핵심으로 사용합니다.

Self-Attention은 시퀀스 내의 모든 위치 쌍에 대해 관련성을 계산합니다.
Query, Key, Value 세 개의 행렬을 사용하며, Scaled Dot-Product Attention으로 계산됩니다.

BERT는 Transformer의 인코더 부분만 사용하는 모델입니다.
양방향 컨텍스트를 활용하여 Masked Language Model과 Next Sentence Prediction으로 사전학습합니다.
텍스트 분류, 개체명 인식, 질의응답 등 다양한 NLP 태스크에서 뛰어난 성능을 보입니다.

GPT는 Transformer의 디코더 부분만 사용하는 모델입니다.
자기회귀 방식으로 다음 토큰을 예측하며, 텍스트 생성에 강점을 가집니다.
GPT-3부터는 In-context Learning 능력이 두드러지며, 별도의 Fine-tuning 없이 프롬프트만으로 다양한 태스크를 수행합니다.

Tokenization은 텍스트를 모델이 처리할 수 있는 단위로 분할하는 과정입니다.
BPE(Byte Pair Encoding), WordPiece, SentencePiece 등의 방법이 있습니다.
서브워드 토큰화는 OOV(Out-of-Vocabulary) 문제를 효과적으로 해결합니다.
"""

# 파일로 저장
with open("sample_nlp_doc.txt", "w", encoding="utf-8") as f:
    f.write(sample_text)

# 문서 로드
documents = load_documents("sample_nlp_doc.txt")
print(f"문서 내용 (앞 200자): {documents[0].page_content[:200]}...")

from langchain.text_splitter import RecursiveCharacterTextSplitter

# 재귀 문자 분할기 - 단락 > 문장 > 단어 순으로 분할 시도
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,         # 청크 최대 크기 (문자 수)
    chunk_overlap=50,       # 청크 간 겹치는 부분 (문맥 유지)
    separators=["\n\n", "\n", ". ", " ", ""],  # 분할 우선순위
    length_function=len,
)

# 문서 분할
chunks = text_splitter.split_documents(documents)

print(f"원본 문서: {len(documents)}개")
print(f"분할된 청크: {len(chunks)}개")
print(f"평균 청크 크기: {sum(len(c.page_content) for c in chunks) / len(chunks):.0f}자")

# 각 청크 확인
for i, chunk in enumerate(chunks):
    print(f"\n--- 청크 {i + 1} ({len(chunk.page_content)}자) ---")
    print(chunk.page_content[:150] + "...")

from langchain_openai import OpenAIEmbeddings

# OpenAI 임베딩 모델 초기화
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",  # 저렴하고 빠른 모델
    # model="text-embedding-3-large",  # 높은 정확도 필요 시
)

# 단일 텍스트 임베딩 테스트
test_embedding = embeddings.embed_query("Transformer는 무엇인가요?")
print(f"임베딩 차원: {len(test_embedding)}")  # 1536차원
print(f"임베딩 샘플: {test_embedding[:5]}...")

# 여러 텍스트 임베딩 (배치)
sample_texts = [
    "Self-Attention 메커니즘",
    "BERT는 양방향 모델입니다",
    "GPT는 자기회귀 모델입니다",
]
batch_embeddings = embeddings.embed_documents(sample_texts)
print(f"배치 임베딩: {len(batch_embeddings)}개 x {len(batch_embeddings[0])}차원")

from langchain_community.vectorstores import Chroma

# ChromaDB에 청크 저장 (임베딩 자동 생성)
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    collection_name="nlp_docs",
    persist_directory="./chroma_db",  # 디스크에 영구 저장
)

print(f"벡터 저장소 생성 완료: {vectorstore._collection.count()}개 벡터")

# 유사도 검색 테스트
query = "BERT는 어떻게 학습하나요?"
results = vectorstore.similarity_search_with_score(query, k=3)

print(f"\n검색 쿼리: '{query}'")
print(f"검색 결과 (상위 {len(results)}개):")
for i, (doc, score) in enumerate(results):
    print(f"\n  [{i + 1}] 유사도 점수: {score:.4f}")
    print(f"      내용: {doc.page_content[:100]}...")

from openai import OpenAI

client = OpenAI()

def answer_question(
    question: str,
    vectorstore: Chroma,
    k: int = 3,
    model: str = "gpt-4o-mini",
) -> dict:
    """문서 기반 질의응답을 수행합니다."""
    # 1. 관련 문서 검색
    results = vectorstore.similarity_search_with_score(question, k=k)
    context_docs = [doc for doc, score in results]
    scores = [score for doc, score in results]

    # 2. 문맥 구성
    context = "\n\n---\n\n".join([doc.page_content for doc in context_docs])

    # 3. LLM에 질문 + 문맥 전달
    response = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "system",
                "content": (
                    "당신은 문서 기반 질의응답 어시스턴트입니다.\n"
                    "주어진 문맥 정보만을 기반으로 질문에 답변하세요.\n"
                    "문맥에 없는 정보는 '제공된 문서에서 해당 정보를 찾을 수 없습니다'라고 답하세요.\n"
                    "답변 끝에 참고한 문맥의 핵심 내용을 간단히 명시하세요."
                )
            },
            {
                "role": "user",
                "content": f"## 문맥 정보\n{context}\n\n## 질문\n{question}"
            }
        ],
        temperature=0.3,  # 사실 기반 답변이므로 낮은 temperature
    )

    answer = response.choices[0].message.content

    return {
        "question": question,
        "answer": answer,
        "sources": [
            {
                "content": doc.page_content[:150] + "...",
                "score": round(score, 4),
            }
            for doc, score in zip(context_docs, scores)
        ],
        "tokens_used": response.usage.total_tokens,
    }

# 테스트
qa_result = answer_question("BERT와 GPT의 차이점은 무엇인가요?", vectorstore)

print(f"질문: {qa_result['question']}")
print(f"\n답변: {qa_result['answer']}")
print(f"\n참고 문서:")
for i, source in enumerate(qa_result["sources"]):
    print(f"  [{i + 1}] (유사도: {source['score']}) {source['content']}")
print(f"\n사용 토큰: {qa_result['tokens_used']}")

from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA

# LLM 초기화
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)

# Retrieval QA 체인 생성
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",  # 검색 결과를 하나의 프롬프트에 합침
    retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
    return_source_documents=True,
)

# 질의응답
result = qa_chain.invoke({"query": "Tokenization이란 무엇인가요?"})
print(f"답변: {result['result']}")
print(f"참고 문서 수: {len(result['source_documents'])}")

import gradio as gr
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from openai import OpenAI
import tempfile
import os

client = OpenAI()
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 전역 벡터 저장소
global_vectorstore = None

def process_file(file) -> str:
    """업로드된 파일을 처리하여 벡터 저장소에 저장합니다."""
    global global_vectorstore

    if file is None:
        return "파일을 업로드해 주세요."

    # 파일 로드
    file_path = file.name
    if file_path.endswith(".pdf"):
        loader = PyPDFLoader(file_path)
    elif file_path.endswith((".txt", ".md")):
        loader = TextLoader(file_path, encoding="utf-8")
    else:
        return "지원하지 않는 파일 형식입니다. PDF 또는 TXT 파일을 업로드하세요."

    documents = loader.load()

    # 청크 분할
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,
        chunk_overlap=50,
    )
    chunks = splitter.split_documents(documents)

    # 벡터 저장소 생성
    global_vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name="uploaded_docs",
    )

    return f"문서 처리 완료: {len(chunks)}개 청크 생성 (원본 {len(documents)}페이지)"

def ask_question(question: str, history: list) -> str:
    """문서 기반 질의응답을 수행합니다."""
    global global_vectorstore

    if global_vectorstore is None:
        return "먼저 문서를 업로드해 주세요."

    if not question.strip():
        return "질문을 입력해 주세요."

    # 관련 청크 검색
    results = global_vectorstore.similarity_search(question, k=3)
    context = "\n\n---\n\n".join([doc.page_content for doc in results])

    # LLM으로 답변 생성
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": (
                    "업로드된 문서의 내용을 기반으로 질문에 답변합니다.\n"
                    "문서에 없는 내용은 답하지 않습니다.\n"
                    "한국어로 답변합니다."
                )
            },
            {
                "role": "user",
                "content": f"문맥:\n{context}\n\n질문: {question}"
            }
        ],
        temperature=0.3,
    )

    answer = response.choices[0].message.content

    # 출처 정보 추가
    sources = "\n\n---\n**참고 문서 청크:**\n"
    for i, doc in enumerate(results):
        sources += f"- 청크 {i + 1}: {doc.page_content[:80]}...\n"

    return answer + sources

# Gradio Blocks UI
with gr.Blocks(
    title="문서 QA 시스템",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown(
        """
        # 문서 QA 시스템
        PDF 또는 텍스트 파일을 업로드한 후 문서 내용에 대해 질문하세요.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="문서 업로드 (PDF / TXT)",
                file_types=[".pdf", ".txt", ".md"],
            )
            process_btn = gr.Button("문서 처리", variant="primary")
            status = gr.Textbox(label="처리 상태", interactive=False)

        with gr.Column(scale=2):
            chatbot = gr.ChatInterface(
                fn=ask_question,
                title="",
                retry_btn="다시 생성",
                clear_btn="대화 초기화",
            )

    # 이벤트 연결
    process_btn.click(
        fn=process_file,
        inputs=file_input,
        outputs=status,
    )

demo.launch(share=False)