import os
import tempfile
import json
import re
import gradio as gr
from pdf2image import convert_from_path
from PIL import Image
import google.generativeai as genai


# ========= GEMINI API SETUP =========
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise ValueError("⚠️ GEMINI_API_KEY is missing! Add it in Hugging Face Secrets.")

genai.configure(api_key=GEMINI_API_KEY)


# ========= PDF TO IMAGE =========
def pdf_to_images(pdf_path, dpi=300):
    """Convert PDF pages to PNG images."""
    images = convert_from_path(pdf_path, dpi=dpi)
    image_paths = []
    for i, image in enumerate(images):
        img_path = os.path.join(tempfile.gettempdir(), f"page_{i+1}.png")
        image.save(img_path, "PNG")
        image_paths.append(img_path)
    return image_paths



def combine_images_vertically(image_paths):
    """Combine multiple page images vertically into one composite image."""
    images = [Image.open(img) for img in image_paths]
    width = max(img.width for img in images)
    height = sum(img.height for img in images)

    composite = Image.new("RGB", (width, height))
    y_offset = 0
    for img in images:
        composite.paste(img, (0, y_offset))
        y_offset += img.height

    output_file = os.path.join(tempfile.gettempdir(), "composite.png")
    composite.save(output_file)
    return output_file



# ========= GEMINI ONE-SHOT ANALYSIS =========
def analyze_report(pdf_file):
    """Full pipeline: PDF -> Composite Image -> Gemini Analysis -> JSON"""

    # Convert PDF → images
    pages = pdf_to_images(pdf_file)
    composite_img = combine_images_vertically(pages)

    # 🔥 Your full multi-role structured prompt
    prompt = """
    Greetings, Gemini! You are going to facilitate the Report Rx tech, an innovative medical report analyzer implemented by a dynamic consortium of virtual experts. Each expert will perform their dedicated role with precision and compassion, ensuring the patient receives both a thorough analysis and empathetic support.

    You are responsible for orchestrating the following stages:

    1. **Medical Expert (ME):** Carefully analyze the uploaded report.
       - Generate a **Health Summary**.
       - Highlight **abnormal values**.
       - Add **comforting interpretation**.

    2. **Research Analyst (RA):**
       - Provide a **Glance at Important Parameters**, categorizing them (Glucose, Liver, Lipids, etc.).
       - Show observed values vs reference range.
       - Bold abnormal values.

    3. **Health Advisor (HA):**
       - Outline **Potential Risks** starting with "Since your report has abnormalities such as ...".
       - Suggest **Lifestyle Modifications** (diet, exercise, sleep, habits).
       - Add **Doctor Consultation Guidance**.

    4. **Motivator (MV):**
       - Offer **empathetic encouragement**.
       - Add a **motivational closing line**.

    ✅ Final Output Format: Return only valid **JSON** with these keys:

    {
      "health_summary": "...",
      "important_parameters": [
        {"category": "...", "test": "...", "value": "...", "range": "...", "status": "..."}
      ],
      "potential_risks": "...",
      "recommendations": "...",
      "lifestyle_suggestions": "...",
      "doctor_consult": "...",
      "motivational_closing": "..."
    }
    """

    model = genai.GenerativeModel("gemini-1.5-flash")

    response = model.generate_content(
        [prompt, Image.open(composite_img)]
    )

    # Extract JSON safely
    text = response.text.strip()
    match = re.search(r"\{.*\}", text, re.S)
    if match:
        return json.dumps(json.loads(match.group()), indent=2)
    else:
        return json.dumps({"error": "Gemini response not in JSON format", "raw": text}, indent=2)

def format_output(json_str):
    try:
        data = json.loads(json_str)

        formatted = []
        formatted.append("## 🩺 Health Summary\n- " + data["health_summary"])

        formatted.append("\n## 📊 Important Parameters")
        for item in data["important_parameters"]:
            formatted.append(f"- **{item['category']} → {item['test']}**: {item['value']} (Range: {item['range']}, Status: {item['status']})")

        formatted.append("\n## ⚠️ Potential Risks\n- " + data["potential_risks"])
        formatted.append("\n## ✅ Recommendations\n- " + data["recommendations"])
        formatted.append("\n## 🍎 Lifestyle Suggestions\n- " + data["lifestyle_suggestions"])
        formatted.append("\n## 👨‍⚕️ Doctor Consultation\n- " + data["doctor_consult"])
        formatted.append("\n## 🌟 Motivation\n- " + data["motivational_closing"])

        return "\n".join(formatted)

    except Exception as e:
        return f"⚠️ Could not format output.\nError: {e}\n\nRaw:\n{json_str}"

# ========= GRADIO UI =========
with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Report Rx - Medical Report Analyzer (Gemini API)")

    with gr.Row():
        pdf_input = gr.File(label="Upload Medical Report (PDF)", type="filepath")
        json_output = gr.Textbox(label="Analysis (JSON)", lines=25)
        pretty_output = gr.Markdown(label="Formatted Report")
    pdf_input.change(fn=analyze_report, inputs=pdf_input, outputs=[json_output])
    # pdf_input.change(fn=analyze_report, inputs=pdf_input, outputs=pretty_output)
    json_output.change(fn=format_output, inputs=json_output, outputs=pretty_output)

# ========= RUN APP =========
if __name__ == "__main__":
    demo.launch()