import os
import requests
import streamlit as st
from PyPDF2 import PdfReader

# Get the Hugging Face API Token from environment variables
HF_API_TOKEN = os.getenv("HF_API_KEY")
if not HF_API_TOKEN:
    raise ValueError("Hugging Face API Token is not set in the environment variables.")

# Hugging Face API URL and header for Gemma 27B-it model
GEMMA_9B_API_URL = "https://api-inference.huggingface.co/models/google/gemma-2-9b-it"

HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"}

def query_model(api_url, payload):
    response = requests.post(api_url, headers=HEADERS, json=payload)
    if response.status_code == 200:
        return response.json()
    else:
        raise ValueError(f"Request failed with status code {response.status_code}: {response.text}")

def extract_pdf_text(uploaded_file):
    pdf_text = ""
    pdf_reader = PdfReader(uploaded_file)
    for page in pdf_reader.pages:
        pdf_text += page.extract_text()
    return pdf_text

# Streamlit app
st.set_page_config(page_title="Gemma 9B-it Chatbot Interface", layout="wide")
st.title("Gemma 9B-it Chatbot Interface")
st.write("Gemma 9B-it Chatbot Interface")

# Initialize session state for conversation and uploaded file
if "conversation" not in st.session_state:
    st.session_state.conversation = []

# File uploader for PDF
uploaded_file = st.file_uploader("Upload a PDF", type="pdf")

# Handle PDF upload and text extraction
if uploaded_file:
    pdf_text = extract_pdf_text(uploaded_file)
    st.write("### PDF Text Extracted:")
    st.write(pdf_text)

# User input for question
question = st.text_input("Ask a question...", "")

# Handle user input and Gemma 27B-it model response
if st.button("Send") and question:
    try:
        with st.spinner("Waiting for the model to respond..."):
            response = query_model(GEMMA_9B_API_URL, {"inputs": question})
            answer = response.get("generated_text", "No response")
            st.write(f"**Gemma 9B-it:** {answer}")
            st.session_state.conversation.append((question, answer))
    except ValueError as e:
        st.error(str(e))

# Custom CSS for chat bubbles
st.markdown(
    """
    <style>
    .chat-container {
        display: flex;
        flex-direction: column;
        gap: 10px;
        margin-top: 20px;
    }
    .user-message {
        align-self: flex-end;
        background-color: #dcf8c6;
        padding: 10px 14px;
        border-radius: 14px;
        max-width: 80%;
    }
    .bot-message {
        align-self: flex-start;
        background-color: #fff;
        padding: 10px 14px;
        border-radius: 14px;
        max-width: 80%;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Display the conversation
if st.session_state.conversation:
    st.write('<div class="chat-container">', unsafe_allow_html=True)
    for user_message, bot_message in st.session_state.conversation:
        st.write(f'<div class="user-message">You: {user_message}</div>', unsafe_allow_html=True)
        st.write(f'<div class="bot-message">Gemma 9B-it: {bot_message}</div>', unsafe_allow_html=True)
    st.write('</div>', unsafe_allow_html=True)