import streamlit as st
import google.generativeai as genai
from PIL import Image
import PyPDF2
import tempfile
import os
from google.api_core import exceptions
from dotenv import load_dotenv
import time
from gtts import gTTS
import base64

load_dotenv()

# Load the API key from the environment variable
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    st.error("Gemini API key not found. Please set the GEMINI_API_KEY environment variable.")
    st.stop()

# Configure the Gemini API
genai.configure(api_key=api_key)
model = genai.GenerativeModel('gemini-1.5-flash')  # Initialize the model

MAX_RETRIES = 3
RETRY_DELAY = 2  # seconds

# Add Chinese (Simplified) to language support
LANGUAGES = {
    "English": "en",
    "Spanish": "es",
    "French": "fr",
    "German": "de",
    "Italian": "it",
    "Portuguese": "pt",
    "Urdu": "ur",
    "Chinese (Simplified)": "zh-cn"
}

def analyze_medical_report(content, content_type, lang):
    prompt = "Analyze this medical report concisely. Provide key findings, diagnoses, and recommendations:"
    
    # Adjust prompt language if not English
    if lang != "en":
        translations = {
            "es": "Analiza este informe médico de manera concisa. Proporcione hallazgos clave, diagnósticos y recomendaciones:",
            "fr": "Analysez ce rapport médical de manière concise. Fournissez les résultats clés, les diagnostics et les recommandations :",
            "de": "Analysieren Sie diesen medizinischen Bericht kurz und prägnant. Geben Sie wichtige Ergebnisse, Diagnosen und Empfehlungen an:",
            "it": "Analizza questo rapporto medico in modo conciso. Fornisci risultati chiave, diagnosi e raccomandazioni:",
            "pt": "Analise este relatório médico de forma concisa. Forneça os principais resultados, diagnósticos e recomendações:",
            "ur": "اس طبی رپورٹ کا مختصر تجزیہ کریں۔ اہم نتائج، تشخیصات، اور سفارشات فراہم کریں:",
            "zh-cn": "简明分析此医疗报告。提供关键发现、诊断和建议:"
        }
        prompt = translations.get(lang, prompt)
    
    for attempt in range(MAX_RETRIES):
        try:
            if content_type == "image":
                response = model.generate_content([prompt, content])
            else:  # text
                response = model.generate_content(f"{prompt}\n\n{content}")
            
            return response.text
        except exceptions.GoogleAPIError as e:
            if attempt < MAX_RETRIES - 1:
                st.warning(f"An error occurred. Retrying in {RETRY_DELAY} seconds... (Attempt {attempt + 1}/{MAX_RETRIES})")
                time.sleep(RETRY_DELAY)
            else:
                st.error(f"Failed to analyze the report after {MAX_RETRIES} attempts. Error: {str(e)}")
                return fallback_analysis(content, content_type)

def generate_tts_audio(text, lang_code):
    # Generate TTS audio from the provided text and language code
    tts = gTTS(text=text, lang=lang_code)
    
    # Save the audio to a temporary file
    audio_path = "audio_output.mp3"
    tts.save(audio_path)
    
    return audio_path

def audio_player(audio_path):
    # Display an audio player in Streamlit
    audio_file = open(audio_path, "rb")
    audio_bytes = audio_file.read()
    st.audio(audio_bytes, format="audio/mp3")

def extract_text_from_pdf(pdf_file):
    # Create a PDF reader object
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    
    # Extract text from each page
    text = ""
    for page_num in range(len(pdf_reader.pages)):
        page = pdf_reader.pages[page_num]
        text += page.extract_text()

    return text

def main():
    st.title("ReportEase AI")
    st.write("AI-driven Medical Report Analyzer with Multilingual Audio Feedback.")
    
    st.write("Upload a Medical report (Image or PDF) for analysis.")

    language = st.selectbox("Select language for analysis and audio feedback:", list(LANGUAGES.keys()))
    lang_code = LANGUAGES[language]

    file_type = st.radio("Select file type:", ("Image", "PDF"))

    if file_type == "Image":
        uploaded_file = st.file_uploader("Choose a medical report image", type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
                tmp_file.write(uploaded_file.getvalue())
                tmp_file_path = tmp_file.name

            image = Image.open(tmp_file_path)
            st.image(image, caption="Uploaded Medical Report", use_column_width=True)

            if st.button("Analyze Image Report"):
                with st.spinner("Analyzing the medical report image..."):
                    analysis = analyze_medical_report(image, "image", lang_code)
                    st.subheader("Analysis Results:")
                    st.write(analysis)

                    # Generate audio of the analysis
                    audio_path = generate_tts_audio(analysis, lang_code)
                    st.write("Listen to the analysis:")
                    audio_player(audio_path)

            os.unlink(tmp_file_path)

    else:  # PDF
        uploaded_file = st.file_uploader("Choose a medical report PDF", type=["pdf"])
        if uploaded_file is not None:
            st.write("PDF uploaded successfully")

            if st.button("Analyze PDF Report"):
                with st.spinner("Analyzing the medical report PDF..."):
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                        tmp_file.write(uploaded_file.getvalue())
                        tmp_file_path = tmp_file.name

                    with open(tmp_file_path, 'rb') as pdf_file:
                        pdf_text = extract_text_from_pdf(pdf_file)

                    analysis = analyze_medical_report(pdf_text, "text", lang_code)
                    st.subheader("Analysis Results:")
                    st.write(analysis)

                    # Generate audio of the analysis
                    audio_path = generate_tts_audio(analysis, lang_code)
                    st.write("Listen to the analysis:")
                    audio_player(audio_path)

                    os.unlink(tmp_file_path)

    # Footer with "Made by Shan"
    st.markdown("---")
    st.markdown("<p style='text-align: center;'>😎 Made by Shan-Ul-Haq 😎</p>", unsafe_allow_html=True)

if __name__ == "__main__":
    main()