"""
앱에 내장된 간단한 RAG 체인 구현
"""
from typing import List, Dict, Any, Optional
import os
from config import OPENAI_API_KEY, LLM_MODEL, USE_OPENAI, TOP_K_RETRIEVAL

# 안전한 임포트
try:
    from langchain_openai import ChatOpenAI
    from langchain.prompts import PromptTemplate
    from langchain_core.output_parsers import StrOutputParser
    from langchain_core.runnables import RunnablePassthrough
    LANGCHAIN_IMPORTS_AVAILABLE = True
except ImportError:
    print("[APP_RAG] langchain 관련 패키지를 로드할 수 없습니다.")
    LANGCHAIN_IMPORTS_AVAILABLE = False

class SimpleRAGChain:
    """
    간단한 RAG 체인 구현 (앱에 내장)
    """
    def __init__(self, vector_store):
        """간단한 RAG 체인 초기화"""
        print("[APP_RAG] 간단한 RAG 체인 초기화 중...")
        self.vector_store = vector_store

        if not LANGCHAIN_IMPORTS_AVAILABLE:
            print("[APP_RAG] langchain 패키지를 찾을 수 없어 RAG 체인을 초기화할 수 없습니다.")
            raise ImportError("RAG 체인 초기화에 필요한 라이브러리가 설치되지 않았습니다.")

        # API 키 확인
        if not OPENAI_API_KEY and USE_OPENAI:
            print("[APP_RAG] 경고: OpenAI API 키가 설정되지 않았습니다.")
            raise ValueError("OpenAI API 키가 설정되지 않았습니다.")

        try:
            # LLM 초기화
            if USE_OPENAI:
                self.llm = ChatOpenAI(
                    model_name=LLM_MODEL,
                    temperature=0.2,
                    api_key=OPENAI_API_KEY,
                )
                print(f"[APP_RAG] OpenAI 모델 초기화: {LLM_MODEL}")
            else:
                try:
                    # Ollama 사용 시도
                    from langchain_community.chat_models import ChatOllama
                    from config import OLLAMA_HOST

                    self.llm = ChatOllama(
                        model=LLM_MODEL,
                        temperature=0.2,
                        base_url=OLLAMA_HOST,
                    )
                    print(f"[APP_RAG] Ollama 모델 초기화: {LLM_MODEL}")
                except ImportError:
                    # Ollama 가져오기 실패 시 OpenAI 사용
                    self.llm = ChatOpenAI(
                        model_name="gpt-3.5-turbo",
                        temperature=0.2,
                        api_key=OPENAI_API_KEY,
                    )
                    print("[APP_RAG] Ollama를 사용할 수 없어 OpenAI로 대체합니다.")

            # 프롬프트 템플릿
            template = """
            다음 정보를 기반으로 질문에 정확하게 답변해주세요.

            질문: {question}

            참고 정보:
            {context}

            참고 정보에 답이 없는 경우 "제공된 문서에서 해당 정보를 찾을 수 없습니다."라고 답변하세요.
            답변은 정확하고 간결하게 제공하되, 참고 정보에서 근거를 찾아 설명해주세요.
            참고 정보의 출처도 함께 알려주세요.
            """

            self.prompt = PromptTemplate.from_template(template)

            # 체인 구성
            self.chain = (
                    {"context": self._retrieve, "question": RunnablePassthrough()}
                    | self.prompt
                    | self.llm
                    | StrOutputParser()
            )
            print("[APP_RAG] RAG 체인 초기화 완료")
        except Exception as e:
            print(f"[APP_RAG] RAG 체인 초기화 실패: {e}")
            import traceback
            traceback.print_exc()
            raise

    def _retrieve(self, query):
        """문서 검색"""
        try:
            docs = self.vector_store.similarity_search(query, k=TOP_K_RETRIEVAL)

            # 검색 결과 컨텍스트 구성
            context_parts = []
            for i, doc in enumerate(docs, 1):
                source = doc.metadata.get("source", "알 수 없는 출처")
                page = doc.metadata.get("page", "")
                source_info = f"{source}"
                if page:
                    source_info += f" (페이지: {page})"

                context_parts.append(f"[참고자료 {i}] - 출처: {source_info}\n{doc.page_content}\n")

            return "\n".join(context_parts)
        except Exception as e:
            print(f"[APP_RAG] 검색 중 오류: {e}")
            import traceback
            traceback.print_exc()
            return "문서 검색 중 오류가 발생했습니다."

    def run(self, query):
        """쿼리 처리"""
        try:
            return self.chain.invoke(query)
        except Exception as e:
            print(f"[APP_RAG] 실행 중 오류: {e}")
            import traceback
            traceback.print_exc()
            return f"오류 발생: {str(e)}"