import os
import smtplib
import gradio as gr
import tensorflow as tf
import numpy as np
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders
from tensorflow.keras.preprocessing import image
from PIL import Image
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.lib.utils import simpleSplit

# Load the trained model
model = tf.keras.models.load_model("my_keras_model.h5")

# Sample images
sample_images = [f"samples/{img}" for img in os.listdir("samples") if img.endswith((".png", ".jpg", ".jpeg"))]

# Email function
def send_email(patient_email, patient_name, pdf_path):
    sender_email = "your_email@gmail.com"
    sender_password = "your_email_password"
    
    try:
        msg = MIMEMultipart()
        msg["From"] = sender_email
        msg["To"] = patient_email
        msg["Subject"] = "Your Bone Fracture Report"

        body = f"Hello {patient_name},\n\nPlease find attached your bone fracture detection report from XYZ Hospital.\n\nBest regards,\nXYZ Hospital"
        msg.attach(MIMEText(body, "plain"))

        with open(pdf_path, "rb") as attachment:
            part = MIMEBase("application", "octet-stream")
            part.set_payload(attachment.read())
            encoders.encode_base64(part)
            part.add_header("Content-Disposition", f"attachment; filename={pdf_path}")
            msg.attach(part)

        server = smtplib.SMTP("smtp.gmail.com", 587)
        server.starttls()
        server.login(sender_email, sender_password)
        server.sendmail(sender_email, patient_email, msg.as_string())
        server.quit()

        return "Report sent successfully!"
    except Exception as e:
        return f"Error sending email: {str(e)}"

# Generate PDF report
def generate_report(name, age, gender, weight, height, allergies, cause, xray, email):
    image_size = (224, 224)

    def predict_fracture(xray_path):
        img = Image.open(xray_path).resize(image_size)
        img_array = image.img_to_array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        prediction = model.predict(img_array)[0][0]
        return prediction

    prediction = predict_fracture(xray)
    diagnosed_class = "Normal" if prediction > 0.5 else "Fractured"
    severity = "Mild" if prediction < 0.3 else "Moderate" if prediction < 0.7 else "Severe"

    # Save X-ray image
    img = Image.open(xray).resize((300, 300))
    img_path = f"{name}_xray.png"
    img.save(img_path)

    # PDF Report
    report_path = f"{name}_fracture_report.pdf"
    c = canvas.Canvas(report_path, pagesize=letter)

    # Header
    c.setFont("Helvetica-Bold", 16)
    c.drawCentredString(300, 770, "XYZ Hospital, New Delhi")
    c.setFont("Helvetica", 12)
    c.drawCentredString(300, 750, "123 Health Street, New Delhi, India")
    c.line(50, 740, 550, 740)

    # Patient Details
    c.setFont("Helvetica-Bold", 14)
    c.drawString(50, 710, "Patient Information:")
    c.setFont("Helvetica", 12)
    details = [
        f"Name: {name}",
        f"Age: {age}",
        f"Gender: {gender}",
        f"Weight: {weight} kg",
        f"Height: {height} cm",
        f"Allergies: {allergies if allergies else 'None'}",
        f"Cause of Injury: {cause if cause else 'Not Provided'}"
    ]
    y = 690
    for detail in details:
        c.drawString(50, y, detail)
        y -= 20

    # Diagnosis
    c.setFont("Helvetica-Bold", 14)
    c.drawString(50, y, "Diagnosis:")
    c.setFont("Helvetica", 12)
    y -= 20
    c.drawString(50, y, f"Fractured: {'Yes' if diagnosed_class == 'Fractured' else 'No'}")
    y -= 20
    c.drawString(50, y, f"Injury Severity: {severity}")

    # X-ray Image
    c.drawInlineImage(img_path, 150, y - 260, width=300, height=300)
    y -= 280

    # Treatment & Recommendations
    c.setFont("Helvetica-Bold", 14)
    c.drawString(50, y, "Recommended Treatment:")
    c.setFont("Helvetica", 12)
    y -= 20
    recommendations = {
        "Mild": "Rest, pain relievers, and follow-up X-ray.",
        "Moderate": "Plaster cast, minor surgery if needed.",
        "Severe": "Major surgery, metal implants, and physiotherapy."
    }
    treatment_text = recommendations[severity]
    for line in simpleSplit(treatment_text, "Helvetica", 12, 480):
        c.drawString(50, y, line)
        y -= 20

    # Estimated Cost
    c.setFont("Helvetica-Bold", 14)
    c.drawString(50, y, "Estimated Treatment Cost:")
    c.setFont("Helvetica", 12)
    y -= 20
    cost_gov = f"Government Hospital: ₹{2000 if severity == 'Mild' else 8000 if severity == 'Moderate' else 20000} - ₹{5000 if severity == 'Mild' else 15000 if severity == 'Moderate' else 50000}"
    cost_priv = f"Private Hospital: ₹{10000 if severity == 'Mild' else 30000 if severity == 'Moderate' else 100000}+"
    for line in simpleSplit(cost_gov, "Helvetica", 12, 480):
        c.drawString(50, y, line)
        y -= 20
    for line in simpleSplit(cost_priv, "Helvetica", 12, 480):
        c.drawString(50, y, line)
        y -= 20

    c.save()

    # Send email with report
    email_status = send_email(email, name, report_path)

    return report_path, email_status

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# Bone Fracture Detection System\n### AI-powered diagnosis and treatment recommendations")

    with gr.Row():
        name = gr.Textbox(label="Patient Name", max_chars=50)
        age = gr.Number(label="Age")

    with gr.Row():
        gender = gr.Radio(["Male", "Female", "Other"], label="Gender")
        email = gr.Textbox(label="Patient Email")

    with gr.Row():
        weight = gr.Number(label="Weight (kg)")
        height = gr.Number(label="Height (cm)")

    with gr.Row():
        allergies = gr.Textbox(label="Allergies (if any)")
        cause = gr.Textbox(label="Cause of Injury (Max 100 words)", max_chars=500)

    with gr.Row():
        xray = gr.Image(type="filepath", label="Upload X-ray Image")

    submit_button = gr.Button("Generate Report")
    output_file = gr.File(label="Download Report")
    email_status = gr.Textbox(label="Email Status", interactive=False)

    submit_button.click(
        generate_report,
        inputs=[name, age, gender, weight, height, allergies, cause, xray, email],
        outputs=[output_file, email_status]
    )

if __name__ == "__main__":
    app.launch()