from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma

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

# LLM 이 '스스로는 모르는' 사내 전용 내용 → 검색이 꼭 필요해진다 (agentic 판단 시연)
sample_docs = [
    Document(page_content="우리 회사의 배포 도구 'DeployX' 는 매주 화요일 새벽 2시에 자동 배포를 수행한다."),
    Document(page_content="DeployX 롤백은 'dx rollback <버전>' 명령으로 약 1분 내에 이전 버전으로 복구한다."),
    Document(page_content="우리 회사의 코드 리뷰 규정: PR 은 최소 2명 승인과 'ready' 라벨이 있어야 머지된다."),
    Document(page_content="우리 회사의 사내 VPN 은 2단계 인증(OTP) 후에만 접속할 수 있다."),
]
splits = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50).split_documents(sample_docs)
store = Chroma.from_documents(splits, embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
                             collection_name="agentic_practice")
retriever = store.as_retriever(search_kwargs={"k": 3})


def needs_retrieval(question: str) -> bool:
    # (1) 정답
    resp = llm.invoke([
        SystemMessage(content=(
            "질문에 답하려면 사내/외부 문서 검색이 필요한가? "
            "특정 회사·제품·도구 이름이나 내부 규정처럼 '일반 상식으로 알 수 없는 고유 정보'가 필요하면 YES. "
            "일반 상식이나 단순 계산이면 NO. YES 또는 NO 한 단어만.")),
        HumanMessage(content=question),
    ])
    return "YES" in resp.content.upper()


def rewrite_query(question: str, prev_query) -> str:
    # (2) 정답
    if prev_query is None:
        prompt = f"다음 질문을 벡터 검색용 키워드 쿼리로 재작성. 쿼리만:\n{question}"
    else:
        prompt = (f"이전 검색이 부족했습니다. 다른 각도로 쿼리 재작성. 쿼리만.\n"
                  f"질문: {question}\n이전 쿼리: {prev_query}")
    return llm.invoke([HumanMessage(content=prompt)]).content.strip()


def is_sufficient(question: str, docs: list[str]) -> bool:
    # (3) 정답
    resp = llm.invoke([
        SystemMessage(content="검색 결과가 질문에 답하기 충분하면 SUFFICIENT, 아니면 INSUFFICIENT. 한 단어만."),
        HumanMessage(content=f"질문: {question}\n\n결과:\n" + "\n".join(docs)),
    ])
    return "SUFFICIENT" in resp.content.upper()


def answer(question: str) -> str:
    # (4) 정답 — 판단/재작성/검색/평가 단계를 print 로 보여 준다
    print(f"[질문] {question}")

    if not needs_retrieval(question):
        print("[판단] 검색 불필요 → 바로 답")
        return llm.invoke([HumanMessage(content=question)]).content

    print("[판단] 검색 필요 → 루프 시작")
    query, documents = None, []
    for attempt in range(1, MAX_TRIES + 1):
        query = rewrite_query(question, query)
        print(f"[재작성 {attempt}] {query}")
        documents = [d.page_content for d in retriever.invoke(query)]
        print(f"[검색] {len(documents)}건 회수")
        if is_sufficient(question, documents):
            print("[평가] 충분 → 루프 종료")
            break
        print("[평가] 부족 → 다른 각도로 재시도")

    context = "\n".join(documents)
    prompt = f"자료만 보고 답하세요. 자료에 없으면 추측 금지.\n\n자료:\n{context}\n\n질문: {question}"
    return llm.invoke([HumanMessage(content=prompt)]).content


if __name__ == "__main__":
    print("[답변]", answer("DeployX 롤백은 어떻게 해?"))
    print()
    print("[답변]", answer("2 + 2 는?"))
    store.delete_collection()
