#!/usr/bin/env python3
"""
Simple RAG System - 첫 번째 RAG 시스템
30분 만에 작동하는 RAG를 만들어보세요!
"""

import os
from typing import List
from langchain.document_loaders import TextLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

class SimpleRAG:
    def __init__(self, persist_directory="./rag_db"):
        """RAG 시스템 초기화"""
        self.persist_directory = persist_directory
        self.embeddings = OpenAIEmbeddings()
        self.llm = ChatOpenAI(
            model_name="gpt-3.5-turbo",
            temperature=0,
            max_tokens=500
        )
        self.vectorstore = None
        self.qa_chain = None
        
    def load_documents(self, file_paths: List[str]):
        """다양한 문서 형식 로드"""
        documents = []
        
        for file_path in file_paths:
            if file_path.endswith('.pdf'):
                loader = PyPDFLoader(file_path)
            elif file_path.endswith('.txt'):
                loader = TextLoader(file_path)
            else:
                print(f"지원하지 않는 파일 형식: {file_path}")
                continue
            
            documents.extend(loader.load())
            print(f"✓ {file_path} 로드 완료")
        
        return documents
    
    def process_documents(self, documents):
        """문서를 청크로 분할"""
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
            separators=["\n\n", "\n", " ", ""]
        )
        
        chunks = text_splitter.split_documents(documents)
        print(f"\n📊 총 {len(chunks)}개의 청크가 생성되었습니다.")
        return chunks
    
    def create_vectorstore(self, chunks):
        """벡터 데이터베이스 생성"""
        print("\n🔄 벡터 데이터베이스 생성 중...")
        
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embeddings,
            persist_directory=self.persist_directory
        )
        
        self.vectorstore.persist()
        print("✓ 벡터 데이터베이스 생성 완료!")
    
    def setup_qa_chain(self):
        """질의응답 체인 설정"""
        # 커스텀 프롬프트 템플릿
        template = """아래의 문맥을 사용하여 질문에 답변해주세요.
        
문맥: {context}

질문: {question}

답변할 때 다음 사항을 지켜주세요:
1. 문맥에 있는 정보만을 사용하여 답변하세요.
2. 문맥에 없는 내용은 추측하지 마세요.
3. 답변을 찾을 수 없다면 "제공된 문서에서 해당 정보를 찾을 수 없습니다"라고 답하세요.

답변:"""
        
        PROMPT = PromptTemplate(
            template=template,
            input_variables=["context", "question"]
        )
        
        # QA 체인 생성
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.vectorstore.as_retriever(
                search_kwargs={"k": 3}
            ),
            return_source_documents=True,
            chain_type_kwargs={"prompt": PROMPT}
        )
        
        print("✓ QA 체인 설정 완료!")
    
    def build(self, file_paths: List[str]):
        """전체 RAG 시스템 구축"""
        print("🚀 RAG 시스템 구축을 시작합니다...\n")
        
        # 1. 문서 로드
        documents = self.load_documents(file_paths)
        
        # 2. 문서 처리
        chunks = self.process_documents(documents)
        
        # 3. 벡터 DB 생성
        self.create_vectorstore(chunks)
        
        # 4. QA 체인 설정
        self.setup_qa_chain()
        
        print("\n✅ RAG 시스템이 준비되었습니다!")
    
    def query(self, question: str):
        """질문에 대한 답변 생성"""
        if not self.qa_chain:
            print("❌ 먼저 build() 메서드를 실행하세요.")
            return None
        
        result = self.qa_chain({"query": question})
        
        return {
            "answer": result["result"],
            "sources": result["source_documents"]
        }

def main():
    # 1. RAG 시스템 초기화
    rag = SimpleRAG()
    
    # 2. 문서 준비 (예시)
    file_paths = [
        "document1.pdf",
        "document2.txt",
    ]
    
    # 3. RAG 시스템 구축
    rag.build(file_paths)
    
    # 4. 대화형 질의응답
    print("\n💬 질문을 입력하세요 (종료: 'quit')\n")
    
    while True:
        question = input("질문> ")
        
        if question.lower() == 'quit':
            print("RAG 시스템을 종료합니다. 안녕히 가세요! 👋")
            break
        
        # 답변 생성
        result = rag.query(question)
        
        if result:
            print(f"\n답변: {result['answer']}")
            print(f"\n참고한 문서 ({len(result['sources'])}개):")
            for i, doc in enumerate(result['sources'], 1):
                print(f"{i}. {doc.page_content[:100]}...")
            print("-" * 50 + "\n")

if __name__ == "__main__":
    # OpenAI API 키 확인
    if not os.getenv("OPENAI_API_KEY"):
        print("⚠️  OPENAI_API_KEY 환경변수를 설정해주세요!")
        print("export OPENAI_API_KEY='your-api-key'")
    else:
        main()