from backend.ai.llm import hugging_llm
from langchain_core.prompts import ChatPromptTemplate
from backend.prompts.all_prompt import SUMMARY_SYSTEM_PROMPT, CALL_ASSISTANT_PROMPT
from backend.repository.models import CallHistory
from backend.schemas.summary_schema import CallSummary, CallCreate
from sqlalchemy.orm import Session
from backend.schemas.assistant_schema import AssistantRequest
from langchain_chroma import Chroma
from backend.ai.embedding import watson_embedding
from langchain_core.output_parsers import StrOutputParser

# 질의 응답
def answer_assistant_question(customer_id:int, question:str, db:Session):
    # 1 단계 : 벡터 DB에서 질의
    # 벡터db 불러오기
    vectorstore = Chroma(embedding_function=watson_embedding, persist_directory="./vectordb")
    # as_retriever()
    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
    # invoke() = docs => page_content join
    docs = retriever.invoke(question)
    sim_context = "\n\n".join(doc.page_content for doc in docs)
    # 2 단계 : DB 검색
    # 고객이 이전에 질문한 내역을 추출
    if customer_id:
        histories = db.query(CallHistory).filter(CallHistory.customer_id == customer_id).order_by(CallHistory.created_at.desc()).limit(5).all()

        # 문제, 해결 컬럼만 문자열로 추출
        customer_text ="\n".join([f"""
        문제: {h.customer_issue}\n
        해결: {h.resolution}
""" for h in histories])

    # 1, 2 단계 => LLM => 답변 생성
    prompt = ChatPromptTemplate.from_template(CALL_ASSISTANT_PROMPT)
    chain = prompt | hugging_llm | StrOutputParser()
    result = chain.invoke({"sim_context": sim_context, "customer_text": customer_text, "question": question})

    return {"answer" : result}
    # return AssistantRequest(answer=result)