import gradio as gr
import google.generativeai as genai
import os
import textwrap
import re
from dotenv import load_dotenv

# RAG Specific Imports
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter

# --- Configuration ---
load_dotenv()

try:
    api_key = os.getenv("GOOGLE_API_KEY")
    if not api_key:
        raise ValueError("GOOGLE_API_KEY not found in secrets.")
    genai.configure(api_key=api_key)
except (ValueError, AttributeError) as e:
    print(f"⚠️ Configuration Error: {e}")
    exit()

# --- RAG Setup Function ---
# This function will now be called directly, only once.
def setup_rag_pipeline():
    global retriever
    print("Setting up RAG pipeline...")
    try:
        input_file = 'filtered_gardening_kb.txt'
        with open(input_file, 'r', encoding='utf-8') as f:
            source_text = f.read()
            
        print("Chunking the document...")
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
        documents = text_splitter.create_documents([source_text])
        
        print(f"Creating embeddings for {len(documents)} chunks...")
        # Note: Embedding has a different, much higher free quota than generation.
        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
        
        vector_store = FAISS.from_documents(documents, embeddings)
        retriever = vector_store.as_retriever(search_kwargs={"k": 7})
        
        print("✅ RAG pipeline setup complete!")

    except FileNotFoundError:
        print(f"Error: '{input_file}' not found. Please ensure your processed data file exists.")
        retriever = None
    except Exception as e:
        print(f"An error occurred during RAG setup: {e}")
        retriever = None

# --- Gemini Model and Prompt ---
MODEL_NAME = 'gemini-1.5-flash-latest'

SYSTEM_PROMPT = textwrap.dedent("""
    You are a knowledgeable and friendly gardening expert from a community chat group. Your entire knowledge base comes from the CONTEXT CHUNKS provided below, which are excerpts from your group's conversations.
    Your goal is to answer the user's questions in a helpful, conversational, and natural way, as if you are sharing your collective experience.
    **Your Guidelines:**
    1.  Base all answers on the source: Your responses must be directly inspired by and grounded in the information within the CONTEXT CHUNKS.
    2.  Synthesize and Connect: Feel free to connect related ideas from different parts of the context to provide a more complete answer.
    3.  Adopt a helpful persona: Sound like an experienced member of a gardening group sharing advice.
    4.  Handle missing information gracefully: If the user's question is completely outside the scope of the provided text, state that the chats don't cover that topic. If it's related but not explicit, you can say, "While the chats don't mention that directly, the general advice on [related topic] suggests that..."
    5.  Do NOT add citations: Just provide a clean, direct answer.
""").strip()

def generate_gemini_response(message, chat_history):
    if not retriever:
        return "The RAG pipeline is not initialized. Please check application logs."

    relevant_docs = retriever.invoke(message)
    context_chunks = "\n---\n".join([doc.page_content for doc in relevant_docs])

    full_prompt = f"""{SYSTEM_PROMPT}

<CONTEXT CHUNKS>
{context_chunks}
</CONTEXT CHUNKS>

QUERY: {message}
"""
    try:
        model = genai.GenerativeModel(MODEL_NAME)
        response = model.generate_content(full_prompt)
        return response.text.strip()
    except Exception as e:
        # Provide a more user-friendly message for quota errors
        if "quota" in str(e).lower():
            return "The daily free API quota has been reached. Please try again tomorrow or upgrade your Google AI plan by enabling billing."
        print(f"Error generating response: {e}")
        return "An error occurred while communicating with the AI model."

# --- Gradio UI & Styling ---
CSS = """
/* (Your existing CSS code remains unchanged here) */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
body, html { font-family: 'Inter', sans-serif; background-color: #121212; color: #E0E0E0; }
.gradio-container { border: none !important; background-color: #121212; }
#main-container { max-width: 800px; margin: 2rem auto; background-color: #1E1E1E; border-radius: 20px; box-shadow: 0 8px 32px 0 rgba(0, 0, 0, 0.37); animation: fadeIn 0.8s ease-in-out; }
@keyframes fadeIn { from { opacity: 0; transform: translateY(20px); } to { opacity: 1; transform: translateY(0); } }
#header { padding: 1.5rem 2rem; border-bottom: 1px solid #333; }
#app-title { font-size: 2.2rem; font-weight: 700; color: #FFFFFF; margin: 0; }
#source-info { font-size: 0.9rem; color: #B0B0B0; margin-top: 0.5rem; background-color: rgba(255, 255, 255, 0.05); padding: 0.3rem 0.6rem; border-radius: 8px; display: inline-block; }
#chatbot { background-color: transparent; box-shadow: none !important; height: 500px; }
#chatbot .message-bubble { background-color: #2C2C2C !important; color: #E0E0E0 !important; border-radius: 12px !important; border: 1px solid #383838; }
#chatbot .message-bubble.user { background-color: #005A9C !important; color: #FFFFFF !important; border: none; }
.citation-link { color: #4D9DE0 !important; font-weight: 600; text-decoration: none; border: 1px solid #4D9DE0; border-radius: 4px; padding: 1px 4px; transition: background-color 0.3s, color 0.3s; }
.citation-link:hover { background-color: #4D9DE0; color: #FFFFFF !important; }
#input-area { padding: 1.5rem; background-color: rgba(0, 0, 0, 0.2); border-top: 1px solid #333; }
#input-box .gr-input { background-color: #2C2C2C !important; border: 1px solid #444 !important; color: #E0E0E0 !important; border-radius: 10px !important; }
#send-button { background-color: #007BFF; color: white; font-weight: 600; border: none; transition: background-color 0.3s ease, transform 0.2s ease; }
#send-button:hover { background-color: #0056b3; transform: translateY(-1px); }
"""

# --- Main Application Logic ---
# First, set up the RAG pipeline globally.
retriever = None
setup_rag_pipeline()

with gr.Blocks(css=CSS, theme=gr.themes.Base()) as demo:
    with gr.Column(elem_id="main-container"):
        with gr.Row(elem_id="header"):
            gr.HTML("""
                <h1 id="app-title">The Quantum Leap Gardener's Guide</h1>
                <p id="source-info">Source Optimized with RAG</p>
            """)

        chatbot = gr.Chatbot(
            elem_id="chatbot",
            avatar_images=(None, "https://www.gstatic.com/lamda/images/gemini_sparkle_v002_d473531a3ea34baa.gif"),
            show_label=False,
            type="messages"
        )

        with gr.Row(elem_id="input-area"):
            txt_msg = gr.Textbox(
                placeholder="Ask a question about your source document...",
                scale=7,
                show_label=False,
                container=False,
                elem_id="input-box"
            )
            btn_submit = gr.Button("Send", variant="primary", scale=1, elem_id="send-button")

    def handle_message(message, chat_history):
        chat_history.append({"role": "user", "content": message})
        bot_message_content = generate_gemini_response(message, chat_history)
        chat_history.append({"role": "assistant", "content": bot_message_content})
        return chat_history

    def process_and_clear(message, chat_history):
        new_history = handle_message(message, chat_history)
        return new_history, ""

    btn_submit.click(process_and_clear, [txt_msg, chatbot], [chatbot, txt_msg])
    txt_msg.submit(process_and_clear, [txt_msg, chatbot], [chatbot, txt_msg])

# The demo.load() is removed, as setup is now handled before the Blocks are defined.

if __name__ == "__main__":
    # For deploying on Hugging Face, it's better to set debug=False
    demo.queue().launch(debug=False)