import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from PyPDF2 import PdfReader
import os

# Load the Gemma model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b-it")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Helper functions for document processing
def get_pdf_text(pdf_file):
    text = ""
    pdf_reader = PdfReader(pdf_file)
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

def load_farming_knowledge_base(pdf_path="ai-farming.pdf"):
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF document '{pdf_path}' not found.")
    return get_pdf_text(pdf_path)

# Load knowledge base from the farming PDF
knowledge_base = load_farming_knowledge_base()

# Chatbot response generation
def chatbot_response(user_message):
    # Check if the question relates to the knowledge base
    if user_message.lower() in knowledge_base.lower():
        context = "This information is extracted from the AI Farming Guide:"
        input_text = f"{context}\n{user_message}\n"
    else:
        context = "Answer based on general farming knowledge:"
        input_text = f"{context}\n{user_message}\n"

    # Generate a response using the Gemma model
    response = pipe(input_text, max_length=512, temperature=0.7, top_p=0.9)
    return response[0]["generated_text"]

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🌾 AI Agri Farmer Chat Bot 🌿")
    gr.Markdown(
        "Welcome to the AI Agri Farmer Chat Bot! 🤖 Ask your farming-related questions, "
        "such as crop management, soil health, fertilizers, or pest control. If your "
        "question isn't found in the farming guide, the bot will answer based on general knowledge."
    )

    with gr.Row():
        user_input = gr.Textbox(
            label="💬 Ask your farming question:",
            placeholder="Example: 'What is the best fertilizer for wheat?'",
        )
        chatbot_output = gr.Textbox(label="🤖 Chat Bot Response:")

    example_inputs = gr.Examples(
        examples=[
            "What is the best fertilizer for rice?",
            "How much water does maize need weekly?",
            "What crops grow well in clay soil?",
        ],
        inputs=user_input,
    )

    def respond(input_text):
        return chatbot_response(input_text)

    user_input.submit(respond, inputs=user_input, outputs=chatbot_output)

    gr.Markdown("### 📄 About the Knowledge Base")
    gr.Markdown(
        "The chatbot uses information from the AI Farming Guide (`ai-farming.pdf`) as its primary source. "
        "For topics not covered, it falls back on general farming knowledge."
    )

# Launch the Gradio app
demo.launch()