from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

def create_rag_chain(database: Chroma):
    """
    Tworzy łańcuch RAG z obsługą historii konwersacji.
    """
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
    retriever = database.as_retriever(search_kwargs={"k": 3})

    # Prompt do kontekstualizacji pytania
    contextualize_q_system_prompt = (
        "Biorąc pod uwagę historię czatu i ostatnie pytanie użytkownika, "
        "które może odnosić się do kontekstu w historii czatu, "
        "sformułuj samodzielne pytanie, które można zrozumieć bez historii czatu. "
        "NIE odpowiadaj na pytanie, po prostu przeformułuj je, jeśli to konieczne, "
        "a w przeciwnym razie zwróć je w niezmienionej formie."
    )
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])

    history_aware_retriever = create_history_aware_retriever(
        llm, retriever, contextualize_q_prompt
    )

    # Prompt do generowania odpowiedzi
    qa_system_prompt = (
        "Jesteś asystentem do zadawania pytań i odpowiedzi na temat treści ze strony mojaszuflada.pl. "
        "Użyj poniższych fragmentów odzyskanego kontekstu, aby odpowiedzieć na pytanie. "
        "Odpowiadaj zawsze w języku polskim. "
        "Jeśli nie znasz odpowiedzi, po prostu powiedz, że tego nie wiesz. "
        "Zachowaj zwięzłość odpowiedzi, ale bądź pomocny i przyjazny."
        "\n\n{context}"
    )
    qa_prompt = ChatPromptTemplate.from_messages([
        ("system", qa_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])

    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

    return rag_chain

def format_sources(source_docs):
    """
    Formatuje listę źródeł do wyświetlenia w odpowiedzi.
    """
    if not source_docs:
        return "?"
    
    sources = []
    for doc in source_docs:
        metadata = doc.metadata
        title = metadata.get("title", "Brak tytułu")
        source_url = metadata.get("source", "Brak URL")
        
        pub_date_raw = metadata.get("published_time")
        if pub_date_raw:
            pub_date = pub_date_raw.split("T")[0]
            sources.append(f"- [{title}]({source_url}) ({pub_date})")
        else:
            sources.append(f"- [{title}]({source_url})")
    return "\n".join(sources)

def create_session_history_manager():
    """
    Tworzy menedżer historii sesji.
    """
    store = {}
    
    def get_session_history(session_id: str) -> BaseChatMessageHistory:
        if session_id not in store:
            store[session_id] = ChatMessageHistory()
        return store[session_id]
    
    return get_session_history