from flask import Flask, render_template, jsonify, request
from src.helper import download_hugging_face_embeddings
from langchain_pinecone import PineconeVectorStore
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv
from src.prompt import *
import os
import traceback


app = Flask(__name__)


load_dotenv(override=True)

PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")

os.environ["PINECONE_API_KEY"] = PINECONE_API_KEY
if GOOGLE_API_KEY:
    os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY


embeddings = download_hugging_face_embeddings()

index_name = os.environ.get("PINECONE_INDEX_NAME", "student-chatbot")
# Embed each chunk and upsert the embeddings into your Pinecone index.
docsearch = PineconeVectorStore.from_existing_index(
    index_name=index_name, embedding=embeddings
)


retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k": 3})

chatModel = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=0,
    max_retries=2,
)
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(chatModel, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)


def build_context_fallback_answer(user_query: str) -> str:
    """Return a best-effort answer using retrieved context only (no LLM call)."""
    try:
        docs = retriever.invoke(user_query)
    except Exception:
        return "Gemini quota is reached, and I could not fetch context right now. Please try again shortly."

    if not docs:
        return "Gemini quota is reached, and I could not find relevant context for this question right now."

    top_doc_text = (docs[0].page_content or "").strip()
    if not top_doc_text:
        return "Gemini quota is reached, but retrieved context is empty. Please try again later."

    answer_line = None
    for line in top_doc_text.splitlines():
        if line.lower().startswith("answer:"):
            answer_line = line.split(":", 1)[1].strip()
            break

    if answer_line:
        return f"Gemini quota reached, so I am answering from stored context: {answer_line}"

    snippet = " ".join(
        part.strip() for part in top_doc_text.splitlines() if part.strip()
    )
    snippet = snippet[:450]
    return "Gemini quota reached, so I am answering from stored context: " f"{snippet}"


@app.route("/")
def index():
    return render_template("chat.html")


@app.route("/get", methods=["GET", "POST"])
def chat():
    msg = request.values.get("msg", "").strip()

    if not msg:
        return "Please enter a question.", 200

    print(msg)

    if not GOOGLE_API_KEY:
        return (
            "GOOGLE_API_KEY is missing. Add it to your .env file and restart the app.",
            200,
        )

    try:
        response = rag_chain.invoke({"input": msg})
        answer = response.get("answer") if isinstance(response, dict) else None
        if not answer:
            return (
                "I could not generate a response right now. Please try rephrasing your question.",
                200,
            )

        print("Response : ", answer)
        return str(answer), 200
    except Exception as e:
        print("Error: ", str(e))
        traceback.print_exc()

        error_text = str(e).lower()
        if (
            "api key" in error_text
            or "permission" in error_text
            or "unauthorized" in error_text
        ):
            return (
                "Your Gemini API key is invalid or missing permissions. Please verify GOOGLE_API_KEY.",
                200,
            )
        if "quota" in error_text or "rate" in error_text or "429" in error_text:
            fallback_answer = build_context_fallback_answer(msg)
            return fallback_answer, 200

        return (
            "I am having trouble reaching the AI service right now. Please try again in a few seconds.",
            200,
        )


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=False)