from qdrant_client import QdrantClient
from rich.console import Console
from langchain_core.documents import Document
from langchain_qdrant import QdrantVectorStore
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_ollama.llms import OllamaLLM
from langchain import hub
from langgraph.graph import START, StateGraph
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.prompts import ChatPromptTemplate


from typing_extensions import List, TypedDict

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage


llm = OllamaLLM(model='llama3.2:latest')


console = Console()


client = QdrantClient(url="localhost", port=6333)
COLLECTION_NAME = "wellness_docs"
embedding_model = "intfloat/e5-base-v2"


embeddings = HuggingFaceEmbeddings(model_name=embedding_model)


vector_store = QdrantVectorStore(
    client=client,
    collection_name=COLLECTION_NAME,
    embedding=embeddings
)


prompt = ChatPromptTemplate([
    ('system',
        """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
        Past History of Messages you can use for reference in case you are asked about previous questions: {history}
        Question: {question} 
        Context: {context} 
        Answer:"""
        ),
    ("user", "{question}")
])


class State(TypedDict):
    question: str
    context: List[Document]
    answer: str
    history: List


retriever = vector_store.as_retriever(search_method='mmr', search_kwargs={"k": 5})
def retrieve(state: State):
    retrieved_docs = retriever.invoke(state['question'])
    return {"context": retrieved_docs}


def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    
    messages = prompt.invoke({
        'question': state['question'],
        'context': docs_content,
        'history': state.get('history', [])
    })

    try:
        response = llm.invoke(messages)
    
    except Exception as e:
        return {'answer': f"Error generating response: {e}", 'history': messages}

    updated_history = state.get('history', []) + [
        HumanMessage(content=state['question']),
        AIMessage(content=response)
    ]

    return  {'answer': response, 'history': updated_history}


graph_builder = StateGraph(State).add_sequence([retrieve, generate])
graph_builder.add_edge(START, "retrieve")

checkpointer = InMemorySaver()
graph = graph_builder.compile(checkpointer=checkpointer)


def get_response(question: str, config):
    result = graph.invoke({
        "question": question,
        },config=config)
    
    return result