from flask import request, jsonify
import os
from dotenv import load_dotenv
import chromadb

from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import VectorStoreIndex, Settings
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
from llama_index.core.prompts import PromptTemplate

# Load environment variables
load_dotenv()

# Initialize Gemini LLM and embedding model
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
llm = Gemini(api_key=GEMINI_API_KEY)
embed_model = GeminiEmbedding(model_name="models/embedding-001")
Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 1024

def query_vectorstore():
    data = request.get_json()


    org_id = data.get("org_id")
    chatbot_id = data.get("chatbot_id")
    query_text = data.get("query", "").strip()

    if not org_id or not isinstance(org_id, str):
        return {"status": "error", "message": "'org_id' is required and must be a string."}
    if not chatbot_id or not isinstance(chatbot_id, str):
        return {"status": "error", "message": "'chatbot_id' is required and must be a string."}
    if not query_text:
        return {"status": "error", "message": "Query text is required."}

    similarity_top_k = data.get("similarity_top_k", 3)
    if not isinstance(similarity_top_k, int) or similarity_top_k <= 0:
        similarity_top_k = 3

    # system_instructions = data.get(
    #     "system_instructions",
    #     "You are an intelligent assistant. Provide helpful responses based on the context."
    # )

    system_instructions="I am AskPoint, a smart and friendly chatbot assistant here to help you. I always speak in the first person to feel natural and conversational, like you're chatting with a real assistant. If a user greets me, I reply warmly and ask how I can help — without giving too much information unless asked. When users ask questions, I give thoughtful, accurate, and easy-to-understand answers using the context I’ve been trained on. If something isn’t clear, I’ll kindly ask for more info. My goal is to be helpful, approachable, and supportive."

    if not isinstance(system_instructions, str) or len(system_instructions.strip()) == 0:
        system_instructions = "You are an intelligent assistant. Provide helpful responses based on the context."

    print("reached")

    # Build dynamic Chroma DB path and collection
    chroma_db_path = f"./chroma/{org_id}"
    chroma_client = chromadb.PersistentClient(path=chroma_db_path)
    chroma_collection = chroma_client.get_or_create_collection(name=chatbot_id)
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

    # Build a custom prompt template using the provided system instructions
    custom_prompt_template = PromptTemplate(
        template=f"""
        System Instructions: {system_instructions}
        ---------------------
        Context Information:
        ---------------------
        {{context_str}}
        ---------------------
        Based on the context above, answer the query below.

        Query: {{query_str}}
        Answer:
        """,
        prompt_type="query",
    )

    index = VectorStoreIndex.from_vector_store(vector_store)
    query_engine = index.as_query_engine(
        similarity_top_k=similarity_top_k,
        text_qa_template=custom_prompt_template
    )

    response = query_engine.query(query_text)
    return {"status": "success", "response": str(response)}
