import os
os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
from flask import Flask, render_template, request, redirect, url_for, session, send_file
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
import tensorflow as tf
import numpy as np
from PIL import Image
import pickle
import io
import matplotlib.pyplot as plt
from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.pdfgen import canvas
from reportlab.lib.units import inch
from datetime import datetime
import logging
from flask_mail import Mail, Message
from flask import jsonify, url_for

app = Flask(__name__)
app.secret_key = "e3f6f40bb8b2471b9f07c4025d845be9"

# Database configuration
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:////tmp/snapsin.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)
migrate = Migrate(app, db)

# Mail configuration
app.config['MAIL_SERVER'] = 'smtp.gmail.com'
app.config['MAIL_PORT'] = 465
app.config['MAIL_USERNAME'] = os.environ.get('MAIL_USERNAME')
app.config['MAIL_PASSWORD'] = os.environ.get('MAIL_PASSWORD')
app.config['MAIL_USE_TLS'] = False
app.config['MAIL_USE_SSL'] = True
mail = Mail(app)

MODEL_PATH = "skin_lesion_model.h5"
HISTORY_PATH = "training_history.pkl"
PLOT_PATH = "/tmp/static/training_plot.png"
LOGO_PATH = "static/logo.jpg"
FORM_TEMPLATE = "form.html"
IMG_SIZE = (224, 224)
CONFIDENCE_THRESHOLD = 0.30

label_map = {
    0: "Melanoma",
    1: "Melanocytic nevus",
    2: "Basal cell carcinoma",
    3: "Actinic keratosis",
    4: "Benign keratosis",
    5: "Dermatofibroma",
    6: "Vascular lesion",
    7: "Squamous cell carcinoma"
}

recommendations = {
    "Melanoma": {
        "solutions": [
            "Consult a dermatologist immediately.",
            "Surgical removal is typically required.",
            "Regular follow-up and screening for metastasis."
        ],
        "medications": ["Interferon alfa-2b", "Vemurafenib", "Dacarbazine"]
    },
    "Melanocytic nevus": {
        "solutions": [
            "Usually benign and requires no treatment.",
            "Monitor for any change in shape or color."
        ],
        "medications": ["No medication necessary unless changes occur."]
    },
    "Basal cell carcinoma": {
        "solutions": [
            "Surgical excision or Mohs surgery.",
            "Topical treatments if superficial.",
            "Radiation in select cases."
        ],
        "medications": ["Imiquimod cream", "Fluorouracil cream", "Vismodegib"]
    },
    "Actinic keratosis": {
        "solutions": [
            "Cryotherapy or topical treatments.",
            "Avoid prolonged sun exposure.",
            "Use of sunscreen regularly."
        ],
        "medications": ["Fluorouracil", "Imiquimod", "Diclofenac gel"]
    },
    "Benign keratosis": {
        "solutions": [
            "Generally harmless and often left untreated.",
            "Can be removed for cosmetic reasons."
        ],
        "medications": ["No medication required unless infected."]
    },
    "Dermatofibroma": {
        "solutions": [
            "Benign skin growth, no treatment needed.",
            "Surgical removal if painful or for cosmetic reasons."
        ],
        "medications": ["No medication needed."]
    },
    "Vascular lesion": {
        "solutions": [
            "Treatment depends on type (e.g., hemangioma).",
            "Laser therapy is commonly used.",
            "Observation if no complications."
        ],
        "medications": ["Beta-blockers (e.g., propranolol for hemangioma)"]
    },
    "Squamous cell carcinoma": {
        "solutions": [
            "Surgical removal is standard.",
            "Follow-up for recurrence or metastasis.",
            "Avoid sun exposure and use sunscreen."
        ],
        "medications": ["Fluorouracil", "Cisplatin", "Imiquimod"]
    },
    "Low confidence": {
        "solutions": [
            "The image is not confidently classified.",
            "Please upload a clearer image or consult a doctor."
        ],
        "medications": ["Not available due to low confidence."]
    },
    "Unknown": {
        "solutions": ["No specific guidance available."],
        "medications": ["N/A"]
    }
}

# Logger
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Database Models
class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(100), nullable=False)
    email = db.Column(db.String(120), unique=True, nullable=False)
    scans = db.relationship('Scan', backref='user', lazy=True)

class Scan(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    patient_name = db.Column(db.String(100), nullable=False)
    patient_gender = db.Column(db.String(20), nullable=False)
    patient_age = db.Column(db.Integer, nullable=False)
    prediction = db.Column(db.String(100), nullable=False)
    confidence = db.Column(db.String(20), nullable=False)
    timestamp = db.Column(db.DateTime, default=datetime.utcnow)
    image_filename = db.Column(db.String(100), nullable=False)

# Load Model
model = None
model_load_error = None
def load_model():
    global model, model_load_error
    try:
        if os.path.exists(MODEL_PATH):
            logger.info("Loading model from %s", MODEL_PATH)
            model = tf.keras.models.load_model(MODEL_PATH, compile=False)
            logger.info("Model loaded successfully")
        else:
            logger.error("Model file %s not found", MODEL_PATH)
            model_load_error = f"Model file {MODEL_PATH} not found"
    except Exception as e:
        logger.error("Failed to load model: %s", str(e))
        model_load_error = f"Model deserialization error: {str(e)}. Please ensure the model is compatible with TensorFlow 2.15.0 or re-save it."

# Attempt to load model at startup
load_model()

# Plot training history
if os.path.exists(HISTORY_PATH):
    try:
        with open(HISTORY_PATH, "rb") as f:
            history_dict = pickle.load(f)
        if "accuracy" in history_dict and "val_accuracy" in history_dict:
            os.makedirs("/tmp/static", exist_ok=True)
            plt.plot(history_dict['accuracy'], label='Train Accuracy')
            plt.plot(history_dict['val_accuracy'], label='Val Accuracy')
            plt.xlabel('Epochs')
            plt.ylabel('Accuracy')
            plt.title('Training History')
            plt.legend()
            plt.grid(True)
            plt.savefig(PLOT_PATH)
            plt.close()
            logger.info("Training plot saved at %s", PLOT_PATH)
        else:
            logger.warning("Invalid training history data in %s", HISTORY_PATH)
    except Exception as e:
        logger.warning("Training history load error: %s", str(e))
else:
    logger.warning("Training history file %s not found", HISTORY_PATH)

def preprocess_image(image_bytes):
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        image = image.resize(IMG_SIZE)
        image_array = tf.keras.utils.img_to_array(image)
        return np.expand_dims(image_array, axis=0) / 255.0
    except Exception as e:
        logger.error("Image preprocessing error: %s", str(e))
        raise

def generate_pdf(report, filepath):
    try:
        c = canvas.Canvas(filepath, pagesize=A4)
        width, height = A4
        y = height - 60

        # Background
        c.setFillColor(colors.Color(0.98, 0.98, 0.99, alpha=1))
        c.rect(0, 0, width, height, fill=1, stroke=0)
        
        # Header background
        c.setFillColor(colors.Color(0.94, 0.96, 0.98, alpha=1))
        c.rect(0, height-120, width, 120, fill=1, stroke=0)

        # Logo
        try:
            if os.path.exists(LOGO_PATH):
                c.setFillColor(colors.white)
                c.rect(65, y-25, 50, 50, fill=1, stroke=1)
                c.setStrokeColor(colors.Color(0.7, 0.7, 0.7, alpha=1))
                c.setLineWidth(1)
                c.rect(65, y-25, 50, 50, fill=0, stroke=1)
                c.drawImage(LOGO_PATH, 67, y-23, width=46, height=46, preserveAspectRatio=True, mask='auto')
            else:
                logger.warning("Logo file %s not found, skipping logo", LOGO_PATH)
        except Exception as e:
            logger.warning("Logo error: %s", str(e))

        # Professional title
        c.setFont("Helvetica-Bold", 22)
        c.setFillColor(colors.Color(0.2, 0.2, 0.2, alpha=1))
        c.drawCentredString(width / 2, y + 5, "Medical Diagnosis Report")
        
        # Subtitle
        c.setFont("Helvetica", 11)
        c.setFillColor(colors.Color(0.5, 0.5, 0.5, alpha=1))
        c.drawCentredString(width / 2, y - 15, "Dermatological Analysis")
        
        # Professional line
        c.setStrokeColor(colors.Color(0.8, 0.8, 0.8, alpha=1))
        c.setLineWidth(1)
        c.line(80, y - 35, width - 80, y - 35)
        
        y -= 80

        def professional_section_box(title, fields, extra_gap=20):
            nonlocal y
            box_height = len(fields) * 20 + 40
            c.setFillColor(colors.Color(0.96, 0.96, 0.96, alpha=0.3))
            c.rect(42, y - box_height - 2, width - 84, box_height, fill=1, stroke=0)
            c.setFillColor(colors.white)
            c.rect(40, y - box_height, width - 80, box_height, fill=1, stroke=1)
            c.setStrokeColor(colors.Color(0.9, 0.9, 0.9, alpha=1))
            c.setFillColor(colors.Color(0.95, 0.95, 0.95, alpha=1))
            c.rect(40, y - 30, width - 80, 30, fill=1, stroke=0)
            c.setFont("Helvetica-Bold", 12)
            c.setFillColor(colors.Color(0.3, 0.3, 0.3, alpha=1))
            c.drawString(55, y - 20, title)
            y -= 45
            c.setFont("Helvetica", 10)
            c.setFillColor(colors.Color(0.2, 0.2, 0.2, alpha=1))
            for label, val in fields.items():
                c.setFont("Helvetica-Bold", 9)
                c.setFillColor(colors.Color(0.4, 0.4, 0.4, alpha=1))
                c.drawString(55, y, f"{label}:")
                c.setFont("Helvetica", 9)
                c.setFillColor(colors.Color(0.2, 0.2, 0.2, alpha=1))
                c.drawString(150, y, str(val))
                y -= 20
            y -= extra_gap

        professional_section_box("Patient Information", {
            "Name": report["name"],
            "Email": report["email"],
            "Gender": report["gender"],
            "Age": f"{report['age']} years"
        })

        confidence_val = float(report["confidence"].replace('%', ''))
        confidence_text = f"{report['confidence']} ({'High' if confidence_val > 85 else 'Moderate' if confidence_val > 70 else 'Low'} Confidence)"
        
        professional_section_box("Diagnostic Results", {
            "Condition": report["prediction"],
            "Confidence": confidence_text,
            "Notes": report["message"] if report["message"] else "No additional notes"
        })

        disease = report["prediction"]
        treatment = recommendations.get(disease, recommendations["Unknown"])

        professional_section_box("Treatment Recommendations", {
            f"{i+1}. {line}": "" for i, line in enumerate(treatment["solutions"])
        })

        professional_section_box("Medication Guidelines", {
            f"{i+1}. {line}": "" for i, line in enumerate(treatment["medications"])
        })

        c.setFillColor(colors.Color(0.98, 0.98, 0.98, alpha=1))
        c.rect(40, 40, width - 80, 70, fill=1, stroke=1)
        c.setStrokeColor(colors.Color(0.9, 0.9, 0.9, alpha=1))
        c.setFont("Helvetica-Bold", 10)
        c.setFillColor(colors.Color(0.4, 0.4, 0.4, alpha=1))
        c.drawString(50, 95, "Medical Disclaimer")
        c.setFont("Helvetica", 8)
        c.setFillColor(colors.Color(0.3, 0.3, 0.3, alpha=1))
        disclaimer_lines = [
            "This report is generated using AI technology for preliminary assessment purposes only.",
            "Results should not replace professional medical consultation and diagnosis.",
            "Please consult a qualified healthcare provider for comprehensive medical evaluation."
        ]
        for i, line in enumerate(disclaimer_lines):
            c.drawString(50, 80 - (i * 10), line)

        c.save()
    except Exception as e:
        logger.error("PDF generation error: %s", str(e))
        raise

@app.route("/")
def home():
    try:
        return redirect(url_for("form"))
    except Exception as e:
        logger.error("Error in home route: %s", str(e))
        return render_template(FORM_TEMPLATE, history_plot=None, result={
            "prediction": "Error",
            "confidence": "N/A",
            "message": f"Failed to load page: {str(e)}",
            "email_status": "N/A"
        })

@app.route("/form")
def form():
    try:
        if not os.path.exists(os.path.join(app.template_folder, FORM_TEMPLATE)):
            logger.error("Template %s not found", FORM_TEMPLATE)
            return jsonify({"error": "Form template not found"}), 500
        if not app.config['MAIL_USERNAME'] or not app.config['MAIL_PASSWORD']:
            logger.warning("Mail configuration missing, email functionality may fail")
        if model_load_error:
            return render_template(FORM_TEMPLATE, history_plot="/training_plot.png", result={
                "prediction": "Error",
                "confidence": "N/A",
                "message": f"Model loading failed: {model_load_error}",
                "email_status": "N/A"
            })
        return render_template(FORM_TEMPLATE, history_plot="/training_plot.png")
    except Exception as e:
        logger.error("Error rendering form: %s", str(e))
        return render_template(FORM_TEMPLATE, history_plot=None, result={
            "prediction": "Error",
            "confidence": "N/A",
            "message": f"Failed to load form: {str(e)}",
            "email_status": "N/A"
        }, status=500)

@app.route("/training_plot.png")
def training_plot():
    try:
        if os.path.exists(PLOT_PATH):
            return send_file(PLOT_PATH, mimetype="image/png")
        else:
            logger.warning("Training plot %s not found", PLOT_PATH)
            return "", 404
    except Exception as e:
        logger.error("Error serving training plot: %s", str(e))
        return "", 500

@app.route("/api/history")
def api_history():
    try:
        user_email = request.args.get('email')
        if not user_email:
            return jsonify({"error": "Email parameter is required"}), 400
        user = User.query.filter_by(email=user_email).first()
        if not user:
            return jsonify([])
        scans = Scan.query.filter_by(user_id=user.id).order_by(Scan.timestamp.desc()).all()
        history_data = [{
            "id": scan.id,
            "prediction": scan.prediction,
            "confidence": scan.confidence,
            "timestamp": scan.timestamp.strftime("%B %d, %Y at %I:%M %p"),
            "patient_name": scan.patient_name,
            "image_url": url_for('uploaded_file', filename=scan.image_filename, _external=True)
        } for scan in scans]
        return jsonify(history_data)
    except Exception as e:
        logger.error("Error in history API: %s", str(e))
        return jsonify({"error": "Internal server error"}), 500

@app.route("/api/email-report/<int:scan_id>")
def email_report(scan_id):
    try:
        scan = Scan.query.get(scan_id)
        if not scan:
            return jsonify({"error": "Report not found"}), 404
        report_data = {
            "name": scan.user.name,
            "email": scan.user.email,
            "gender": scan.patient_gender,
            "age": scan.patient_age,
            "prediction": scan.prediction,
            "confidence": scan.confidence,
            "message": ""
        }
        pdf_path = f"/tmp/report_{scan_id}.pdf"
        generate_pdf(report_data, pdf_path)
        msg = Message(
            'Your SnapSkin Diagnostic Report',
            sender=app.config['MAIL_USERNAME'],
            recipients=[scan.user.email]
        )
        msg.body = f"Dear {scan.user.name},\n\nPlease find your requested diagnostic report attached.\n\nThank you for using SnapSkin."
        with app.open_resource(pdf_path) as fp:
            msg.attach(f"SnapSkin_Report_{scan_id}.pdf", "application/pdf", fp.read())
        mail.send(msg)
        os.remove(pdf_path)
        return jsonify({"success": True, "message": f"Report sent to {scan.user.email}"})
    except Exception as e:
        logger.error(f"Failed to send email for scan {scan_id}: {e}")
        return jsonify({"success": False, "message": "Failed to send email."}), 500

@app.route("/predict", methods=["POST"])
def predict():
    try:
        if model_load_error or not model:
            raise ValueError(f"Model not loaded: {model_load_error}")
        if "image" not in request.files:
            raise ValueError("No image uploaded.")
        image = request.files["image"]
        image_bytes = image.read()
        img_array = preprocess_image(image_bytes)
        prediction = model.predict(img_array)[0]
        predicted_index = int(np.argmax(prediction))
        confidence = float(prediction[predicted_index])
        label = label_map.get(predicted_index, "Unknown") if confidence >= CONFIDENCE_THRESHOLD else "Low confidence"
        msg = "⚠ This image is not confidently recognized. Please upload a clearer image." if confidence < CONFIDENCE_THRESHOLD else ""
        email = request.form.get("email")
        user = User.query.filter_by(email=email).first()
        if not user:
            user = User(name=request.form.get("name"), email=email)
            db.session.add(user)
            db.session.commit()
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        image_filename = f"scan_{timestamp}.jpg"
        image_path = os.path.join("static/uploads", image_filename)
        os.makedirs("static/uploads", exist_ok=True)
        image.seek(0)
        image.save(image_path)
        scan = Scan(
            user_id=user.id,
            patient_name=request.form.get("name"),
            patient_gender=request.form.get("gender"),
            patient_age=int(request.form.get("age")),
            prediction=label,
            confidence=f"{confidence * 100:.2f}%",
            image_filename=image_filename
        )
        db.session.add(scan)
        db.session.commit()
        report = {
            "name": request.form.get("name"),
            "email": email,
            "gender": request.form.get("gender"),
            "age": request.form.get("age"),
            "prediction": label,
            "confidence": f"{confidence * 100:.2f}%",
            "message": msg,
            "scan_id": scan.id
        }
        session["report"] = report
        try:
            if not app.config['MAIL_USERNAME'] or not app.config['MAIL_PASSWORD']:
                raise ValueError("Mail configuration missing")
            pdf_path = f"/tmp/report_{scan.id}.pdf"
            generate_pdf(report, pdf_path)
            msg = Message(
                'Your SnapSkin Diagnostic Report',
                sender=app.config['MAIL_USERNAME'],
                recipients=[email]
            )
            msg.body = f"Dear {report['name']},\n\nPlease find your diagnostic report attached.\n\nThank you for using SnapSkin."
            with app.open_resource(pdf_path) as fp:
                msg.attach(f"SnapSkin_Report_{scan.id}.pdf", "application/pdf", fp.read())
            mail.send(msg)
            os.remove(pdf_path)
            report["email_status"] = "Report sent to your email."
        except Exception as e:
            logger.error(f"Failed to send email: {e}")
            report["email_status"] = "Failed to send report to email."
        return redirect(url_for("result"))
    except Exception as e:
        logger.error("Prediction error: %s", str(e))
        return render_template(FORM_TEMPLATE, history_plot="/training_plot.png", result={
            "prediction": "Error",
            "confidence": "N/A",
            "message": f"Prediction failed: {str(e)}",
            "email_status": "Error occurred, no email sent."
        })

@app.route("/result")
def result():
    try:
        if not os.path.exists(os.path.join(app.template_folder, FORM_TEMPLATE)):
            logger.error("Template %s not found", FORM_TEMPLATE)
            return jsonify({"error": "Form template not found"}), 500
        report = session.get("report", {})
        return render_template(FORM_TEMPLATE, **report)
    except Exception as e:
        logger.error("Error rendering result: %s", str(e))
        return render_template(FORM_TEMPLATE, history_plot="/training_plot.png", result={
            "prediction": "Error",
            "confidence": "N/A",
            "message": f"Failed to load result: {str(e)}",
            "email_status": "N/A"
        }, status=500)

@app.route("/download-report")
def download_report():
    try:
        report = session.get("report", {})
        if not report:
            return redirect(url_for("form"))
        os.makedirs("/tmp/reports", exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        filepath = f"/tmp/reports/report_{timestamp}.pdf"
        generate_pdf(report, filepath)
        return send_file(filepath, as_attachment=True)
    except Exception as e:
        logger.error("Download report error: %s", str(e))
        return redirect(url_for("form"))

@app.route("/uploads/<filename>")
def uploaded_file(filename):
    try:
        file_path = os.path.join("static/uploads", filename)
        if os.path.exists(file_path):
            return send_file(file_path)
        else:
            logger.warning("Image file %s not found", file_path)
            return "", 404
    except Exception as e:
        logger.error("Error serving uploaded file: %s", str(e))
        return "", 500

if __name__ == "__main__":
    try:
        with app.app_context():
            db.create_all()
        static_files = ["form-styles.css", "preloader.js", "cursor-effect.js", "logo.png"]
        for file in static_files:
            if not os.path.exists(os.path.join("static", file)):
                logger.warning("Static file %s not found", file)
        app.run(host="0.0.0.0", port=7860)
    except Exception as e:
        logger.error("Application startup error: %s", str(e))
        raise