import gradio as gr
import openai
from PyPDF2 import PdfReader   # Updated import
import io
from openai import OpenAI 
import os

# Define your OpenAI API key
openai.api_key = os.environ.get('openai_api_key')

#client = openai.OpenAI(api_key="sk-none")

client = openai.OpenAI(api_key=os.environ.get('openai_api_key'))

# Function to read PDF and extract text
def pdf_to_text(pdf_file):
    try:
        pdf_reader = PdfReader(io.BytesIO(pdf_file))
        text_content = ''
        for page in pdf_reader.pages:
            text_content += page.extract_text()
        return text_content
    except Exception as e:
        return f"Failed to extract text from PDF: {str(e)}"

def call_openai_model(text_content, question):
    try:
        conversation = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": text_content},
            {"role": "user", "content": question}
        ]
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=conversation
        )
        #return response["choices"][0]["message"]["content"]
        return response.choices[0].message.content
    except Exception as e:
        return f"OpenAI API error: {str(e)}"

# Define a stateful function to keep the conversation history
history = []

def stateful_chatbot(pdf_file=None, question=""):
    try:
        if pdf_file is not None:
            text_content = pdf_to_text(pdf_file)
            history.append(('Document', text_content))
        else:
            text_content = history[-1][1] if history else ""

        if question:
            answer = call_openai_model(text_content, question)
            history.append(('Q', question))
            history.append(('A', answer))
        else:
            answer = ""

        conversation = "\n".join([f"{turn[0]}: {turn[1]}" for turn in history])
    except Exception as e:
        conversation = "An error occurred: " + str(e)
        answer = "An error occurred: " + str(e)
    return  answer

    # Gradio interface with error catching
iface = gr.Interface(
    fn=stateful_chatbot,
    inputs=[
        gr.File(label="Upload PDF Document", type="binary"),
        gr.Textbox(lines=2, placeholder="Ask a question...")
    ],
    outputs=[
        #gr.Textbox(label="Conversation History")
        gr.Textbox(label="Answer")
    ],
    allow_flagging="never",
    live=False,
    theme=gr.themes.Monochrome(
                             primary_hue="red",
                             secondary_hue="red",
                             neutral_hue="red"),
    #title="INSTRUCTIONS:",
             description="This is a helpful AI assistant that allows you, a CKD or dialysis patient, to upload your lab report or bloodwork results in PDF format, the report you get by scanning the QR code (with phone camera or webcam) at the top right of the printout given to you by the hospital. After downloading that PDF report to your PC, mobile phone or tablet, upload it to this chatbot and ask questions about those results and get answers from a kind, friendly, and professional doctor. Be sure to compare these answers with what your real, personal or family Doctor offers. The chatbot will try to find the answer from the lab report and share it with you. You can ask questions about any aspect of your lab results, such as what they mean, how they compare to normal ranges, what causes them, or what actions you should take. Again, the chatbot is not a substitute for professional medical advice, diagnosis, or treatment. Always seek the advice of your physician or other qualified health provider with any questions you may have regarding a medical condition.<br><br><strong>INSTRUCTIONS:</strong><br><br>1. Upload your lab report in PDF format (If you have a photo on your phone instead of a PDF, go to https://ckd-ai.com/images). If results are inaccurate or null, ensure the document uploaded is not password protected, is editable text, and is not just a scanned image of a document.\n\n2. Type your question in the text box below. The question should be relevant to and potentially answerable in the document you uploaded.\n\n3. Click ‘Submit’ to start the answering process. An answer is returned shortly after in the scrollable output box, depending on the complexity of the question.\n\n4. Click ‘Clear’ to remove the submitted document, question, and answer so that you can upload another document and ask another question.")



iface.launch()