import openai
import gradio as gr
import os
import PyPDF2
from docx import Document
from PIL import Image
import pytesseract
import json

# Load your OpenAI API key and prompt from environment variables
openai.api_key = os.getenv("HF_API_KEY_OPENAI")
preset_prompt = os.getenv("HF_PROMPT_HEART_HEALTH")

# Risk Disclaimer to be added at the end of the report
risk_disclaimer = f'''
இந்த AI உருவாக்கிய சுருக்க அறிக்கை, முடிவு ஆதரவு கருவியாக மட்டுமே பயன்பட வேண்டியது ஆகும் மற்றும் மருத்துவ பரிசோதனை எனக் கருதப்படக்கூடாது. இந்த அறிக்கைகள், தனிநபர்களுக்கும் சுகாதார நிபுணர்களுக்கும் உதவிட உருவாக்கப்பட்டவை, மற்றும் க்ளினிக்கல் தீர்மானம் மற்றும் பிற பரிசோதனை முறைகளுடன் இணைந்து பயன்படுத்தப்பட வேண்டும். AI அமைப்பு, தகுதியான மருத்துவ நிபுணர்களின் நிபுணத்துவத்தை மாற்றுவதில்லை. AI உருவாக்கிய அறிக்கைகளின் அடிப்படையில் மட்டும், உரிய மருத்துவ நிபுணருடன் ஆலோசனை இன்றி, எந்த சுகாதார முடிவுகளும் எடுக்கக்கூடாது.
'''

# Function to extract text from a PDF
def extract_text_from_pdf(pdf_file):
    reader = PyPDF2.PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    return text

# Function to extract text from a DOCX file
def extract_text_from_docx(docx_file):
    doc = Document(docx_file)
    text = ""
    for paragraph in doc.paragraphs:
        text += paragraph.text
    return text

# Function to extract text from an image using OCR
def extract_text_from_image(image_file):
    image = Image.open(image_file)
    text = pytesseract.image_to_string(image)
    return text

# Function to detect file type and extract text accordingly
def extract_text(file):
    file_type = file.name.split('.')[-1].lower()
    if file_type == 'pdf':
        return extract_text_from_pdf(file)
    elif file_type == 'docx':
        return extract_text_from_docx(file)
    elif file_type in ['jpg', 'jpeg', 'png']:
        return extract_text_from_image(file)
    elif file_type == 'txt':
        return file.read().decode('utf-8')
    else:
        return "Unsupported file format."

# Function to interact with the GPT-4o-mini model via the OpenAI API
def heart_health_bot(file):
    if file is None:
        return "இதய ஆரோக்கிய அறிக்கையைப் பெற ஒரு ஆவணத்தை (PDF, DOCX, படம், அல்லது உரை) பதிவேற்றவும்."

    context = extract_text(file)
    if "Unsupported file format" in context:
        return context
    
    # Use the secret prompt directly in the API call
    response = openai.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": preset_prompt},
            {"role": "user", "content": context}
        ],
        max_tokens=3000,  # Increased to capture the complete analysis
        temperature=0.7,
    )
    
    # Extract the response text and remove any leading/trailing whitespace
    result = response.choices[0].message.content.strip()

    # Append the risk disclaimer to the result
    result_with_disclaimer = f"{result}\n\n{risk_disclaimer}"

    # Create a JSON output
    output_json = {
        "result": result_with_disclaimer
    }
    
    # Save the JSON output to a file
    json_file = "heart_health_report.json"
    with open(json_file, "w") as f:
        json.dump(output_json, f)

    return result_with_disclaimer, json_file

# Create the Gradio interface
iface = gr.Interface(
    fn=heart_health_bot,
    inputs=[
        gr.File(label="Upload a PDF, DOCX, image, or text file for analysis")  # File upload input
    ],
    outputs=[
        "text",  # Output type for the bot's response
        gr.File(label="Download JSON Report")  # JSON download link
    ],
    title="HeartHealthAI Bot",
    description=(
        "உங்கள் இரத்த பரிசோதனை அல்லது எக்கோ அல்லது ECG அல்லது உங்களின் சமீபத்திய மருத்துவ அறிக்கையை பதிவேற்றவும், இதய ஆரோக்கியத்தை மேம்படுத்துவதற்கான பார்வையைப் பெறுங்கள்..\n\n"
        "Developed by Venkat Srinivasan.\n"
        "Email: venkat.srinivasan@gbs-alumni.com\n"
        "LinkedIn: [Venkataraghavan Srinivasan](https://www.linkedin.com/in/venkataraghavansrinivasan/)\n\n"
        "இதய ஆரோக்கிய அறிக்கையைப் பெற ஒரு ஆவணத்தை (PDF, DOCX, படம், அல்லது உரை) பதிவேற்றவும்.."
    )
)

# Launch the Gradio interface
iface.launch()