import gradio as gr
from huggingface_hub import InferenceClient
from PyPDF2 import PdfReader

# Initialize the Inference Client
client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    uploaded_pdf=None
):
    messages = [{"role": "system", "content": system_message}]

    # Add previous conversation history to the messages
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # If a new message is entered, add it to the conversation history
    messages.append({"role": "user", "content": message})

    # If a PDF is uploaded, process its content
    if uploaded_pdf is not None:
        file_content = extract_pdf_text(uploaded_pdf)
        if file_content:
            messages.append({"role": "user", "content": f"Document Content: {file_content}"})

    # Get response from the model
    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response


def extract_pdf_text(file):
    """Extract text from a PDF file."""
    try:
        reader = PdfReader(file)
        text = ""
        for page in reader.pages:
            text += page.extract_text()
        return text
    except Exception as e:
        return f"Error extracting text from PDF: {str(e)}"


# CSS for styling the interface
css = """
body {
    background-color: #1e2a38; /* Dark blue background */
    color: #ffffff; /* White text for readability */
    font-family: 'Arial', sans-serif; /* Clean and modern font */
}
.gr-button {
    background-color: #42B3CE !important; /* Light blue button */
    color: #2e3b4e !important; /* Dark text for contrast */
    border: none !important;
    padding: 10px 20px !important;
    border-radius: 8px !important;
    font-size: 16px;
    font-weight: bold;
    transition: background-color 0.3s ease, transform 0.2s ease;
}
.gr-button:hover {
    background-color: #3189A2 !important; /* Darker blue on hover */
    transform: scale(1.05);
}
.gr-button:active {
    background-color: #267b88 !important; /* Even darker when clicked */
}
.gr-slider-container {
    color: white !important; /* White slider labels */
    font-size: 14px;
}
.gr-slider {
    background-color: #0b2545 !important; /* Slider track color */
    border-radius: 8px;
}
.gr-slider .gr-slider-active {
    background-color: #42B3CE !important; /* Active slider color */
}
.gr-textbox input {
    background-color: #2f3b4d;
    color: white;
    border: 2px solid #42B3CE;
    padding: 12px;
    border-radius: 8px;
    font-size: 16px;
    transition: border 0.3s ease;
}
.gr-textbox input:focus {
    border-color: #3189A2;
}
"""

# Gradio interface
demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox(label="Your Message", placeholder="Type your question here...", lines=4),
        gr.File(label="Upload a PDF", file_count="single", type="file"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens", visible=False),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature", visible=False),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)", visible=False),
    ],
    outputs="text",
    css=css,  # Custom CSS
    live=True,
    title="Health Assistant Chat",
    description="This is a health assistant that can chat with you about health-related topics. You can also upload a document for analysis.",
)

if __name__ == "__main__":
    demo.launch(share=True)