__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

import os
from dotenv import load_dotenv

# ✅ Load .env FIRST before any other import
load_dotenv()

# ✅ Now .env values are available — no need to hardcode os.environ
# ANONYMIZED_TELEMETRY=False and CHROMA_TELEMETRY=False come from .env

# Patch ChromaDB telemetry bug
import chromadb.telemetry.product.posthog as posthog_module
import unittest.mock as mock
posthog_module.Posthog = mock.MagicMock()

import warnings
warnings.filterwarnings("ignore", category=FutureWarning)

from langchain.prompts import PromptTemplate
from langchain.retrievers import MergerRetriever
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA

# Initialize the embedding model
embeddings = HuggingFaceEmbeddings(
    model_name="all-MiniLM-L6-v2"
)


pdf_vectorstore = Chroma(
    persist_directory="chroma_db",
    collection_name="langchain",
    embedding_function=embeddings
)

video_vectorstore = Chroma(
    persist_directory="chroma_db",
    collection_name="video_chunks",
    embedding_function=embeddings
)

pdf_retriever = pdf_vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 3}
)
video_retriever = video_vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 3}
)

merged_retriever = MergerRetriever(
    retrievers=[pdf_retriever, video_retriever]
)

llm = ChatGroq(
    model="llama-3.1-8b-instant",
    api_key=os.getenv("GROQ_API_KEY"),
    max_tokens=512
)

prompt_template = """You are a helpful AI assistant.

Your task is to read the provided context and answer the question in a clear, meaningful, and human-friendly way.

Instructions:
- Understand the context first, then explain the answer in your own words.
- Do NOT copy sentences directly from the context.
- Use only the information available in the context.
- If the context has partial information, use it to form the best possible complete answer.
- Keep the answer clear, concise, and grammatically correct.
- Ensure the answer is a complete sentence with proper meaning.

If no relevant information is found in the context, respond with:
"I don't have enough information to answer that."


Context:
{context}

Question: {question}

Answer:"""

prompt = PromptTemplate(
    input_variables=["context", "question"],
    template=prompt_template
)

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff", 
    retriever=merged_retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt}
)

def chat():

    while True:
        query = input("You: ").strip()
        if query.lower() in ["exit", "quit"]:
            print("Goodbye!")
            break
        if not query:
            continue

        print(f"  Question : {query}")

        result = qa_chain.invoke({"query": query})

        print(f"\n💬 Answer: {result['result']}")
    

if __name__ == "__main__":
    chat()