from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

confidence_llm = init_chat_model("gpt-4o-mini", temperature=0)

def check_confidence(state: FlareState) -> FlareState:
    """생성된 문장의 확신도를 평가합니다."""
    question = state["question"]
    sentence = state["current_sentence"]

    # 원 FLARE 논문은 토큰별 logprob 확률로 확신도를 판단합니다.
    # 아래는 FLARE-inspired 구현으로, API 제한 환경에서
    # LLM 자기 평가로 대체한 간소화 버전입니다.
    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "다음 문장이 질문에 대해 사실에 기반한 확신 있는 답변인지 평가하세요.\n"
         "추측이나 불확실한 내용이 포함되어 있으면 'low'를 출력하세요.\n"
         "확신 있는 사실적 내용이면 'high'를 출력하세요.\n"
         "'low' 또는 'high'만 출력하세요."),
        ("human", "질문: {question}\n문장: {sentence}"),
    ])

    chain = prompt | confidence_llm | StrOutputParser()
    result = chain.invoke({"question": question, "sentence": sentence})
    is_low = "low" in result.strip().lower()

    return {"is_low_confidence": is_low}

from langchain_chroma import Chroma

vectorstore = Chroma(collection_name="documents", embedding_function=embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

def active_retrieve(state: FlareState) -> FlareState:
    """불확실한 문장을 쿼리로 사용하여 검색합니다."""
    # 임시 생성 문장 자체를 검색 쿼리로 활용
    query = state["current_sentence"]
    documents = retriever.invoke(query)
    return {"documents": documents}

def regenerate_sentence(state: FlareState) -> FlareState:
    """검색 결과를 참고하여 문장을 재생성합니다."""
    question = state["question"]
    partial = state.get("partial_answer", "")
    docs = state.get("documents", [])

    context = "\n\n".join(doc.page_content for doc in docs)

    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "다음 컨텍스트를 참고하여 질문에 대한 답변을 이어서 작성하세요. "
         "정확히 1~2문장만 생성하세요. "
         "컨텍스트에 근거한 사실만 작성하세요.\n\n"
         "컨텍스트:\n{context}"),
        ("human",
         "질문: {question}\n\n"
         "현재까지의 답변: {partial}\n\n"
         "다음 1~2문장을 이어서 작성하세요."),
    ])

    chain = prompt | llm | StrOutputParser()
    sentence = chain.invoke({
        "context": context,
        "question": question,
        "partial": partial,
    })

    return {"current_sentence": sentence.strip()}

def confirm_sentence(state: FlareState) -> FlareState:
    """확정된 문장을 부분 답변에 추가합니다."""
    partial = state.get("partial_answer", "")
    sentence = state["current_sentence"]

    if "[완료]" in sentence:
        return {"generation": partial.strip()}

    updated = f"{partial} {sentence}".strip()
    return {"partial_answer": updated}