import os
import uuid
import json
import numpy as np
import cv2
from PIL import Image
from flask import Flask, render_template, request, redirect, url_for, session, flash
from werkzeug.security import generate_password_hash, check_password_hash
import sqlite3

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

import tensorflow as tf
keras = tf.keras

# Monkeypatch Keras 3 Layer.__init__ to ignore unrecognized quantization_config argument
# which is present in legacy models but unsupported by newer Keras 3 versions.
original_layer_init = keras.layers.Layer.__init__
def patched_layer_init(self, *args, **kwargs):
    kwargs.pop('quantization_config', None)
    original_layer_init(self, *args, **kwargs)
keras.layers.Layer.__init__ = patched_layer_init

load_model = keras.models.load_model

# ──────────────────────────────────────────────
# Flask App Setup
# ──────────────────────────────────────────────

app = Flask(__name__)
app.secret_key = os.environ.get("SECRET_KEY", "ojos_ai_secret_2024_secure")

UPLOAD_FOLDER = os.path.join("static", "uploads")
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config["MAX_CONTENT_LENGTH"] = 10 * 1024 * 1024  # 10 MB max upload

# ──────────────────────────────────────────────
# Load Model
# ──────────────────────────────────────────────

MODEL_PATH   = "best_eye_disease_model.h5"
INDICES_PATH = "class_indices.json"

print("Loading OJOS AI model (96.69% accuracy)...")
eye_model = load_model(MODEL_PATH, compile=False)

with open(INDICES_PATH, "r") as f:
    class_to_idx = json.load(f)

# Invert: idx → class_name
CLASS_NAMES = {v: k for k, v in class_to_idx.items()}
print(f"Model loaded! Classes: {CLASS_NAMES}")

# ──────────────────────────────────────────────
# Disease Info Dictionary
# ──────────────────────────────────────────────

DISEASE_INFO = {
    "cataract": {
        "display": "Cataract Detected",
        "emoji": "👁️",
        "advice": "Cataracts cause clouding of the eye's lens leading to blurry vision. Early surgical treatment (phacoemulsification) is highly effective. Please consult an ophthalmologist at the earliest.",
        "severity": "warning",
        "color": "#f59e0b",
        "tips": [
            "Avoid prolonged exposure to UV light",
            "Schedule a comprehensive eye exam",
            "Surgery is safe and highly effective",
            "Wear UV-protective sunglasses outdoors"
        ]
    },
    "diabetic_retinopathy": {
        "display": "Diabetic Retinopathy",
        "emoji": "⚠️",
        "advice": "Diabetic retinopathy damages blood vessels in the retina due to high blood sugar. Strict blood sugar control and regular eye exams are critical. Consult both an endocrinologist and a retinal specialist immediately.",
        "severity": "danger",
        "color": "#ef4444",
        "tips": [
            "Control blood sugar levels strictly",
            "Get retinal exams every 6 months",
            "Maintain healthy blood pressure",
            "Laser treatment may be required"
        ]
    },
    "glaucoma": {
        "display": "Glaucoma Detected",
        "emoji": "🔴",
        "advice": "Glaucoma damages the optic nerve, often due to elevated eye pressure. It can lead to permanent vision loss if untreated. See an ophthalmologist immediately for pressure testing and treatment.",
        "severity": "danger",
        "color": "#ef4444",
        "tips": [
            "Get intraocular pressure checked immediately",
            "Eye drops or surgery may be needed",
            "Regular monitoring is essential",
            "Early detection prevents blindness"
        ]
    },
    "healthy": {
        "display": "Healthy Eye",
        "emoji": "✅",
        "advice": "Your retinal image appears normal with no signs of eye disease. Continue annual eye checkups and maintain a healthy lifestyle to preserve your vision for years to come.",
        "severity": "success",
        "color": "#10b981",
        "tips": [
            "Continue annual eye checkups",
            "Maintain a healthy diet rich in vitamins",
            "Limit screen time and take breaks",
            "Protect eyes from UV radiation"
        ]
    },
    "myopia": {
        "display": "Myopia Detected",
        "emoji": "👓",
        "advice": "Myopia (nearsightedness) causes difficulty seeing distant objects clearly. Corrective lenses, contact lenses, or LASIK surgery can restore clear vision. Consult an optometrist for the best treatment option.",
        "severity": "info",
        "color": "#3b82f6",
        "tips": [
            "Get prescription glasses or contacts",
            "LASIK surgery is an effective option",
            "Take regular breaks from screens",
            "Spend more time outdoors — proven to slow progression"
        ]
    }
}

# ──────────────────────────────────────────────
# Database Setup
# ──────────────────────────────────────────────

def get_db():
    db_path = "/tmp/users.db" if os.name != "nt" else "users.db"
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    conn = get_db()
    conn.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id       INTEGER PRIMARY KEY AUTOINCREMENT,
            name     TEXT    NOT NULL,
            email    TEXT    UNIQUE NOT NULL,
            password TEXT    NOT NULL
        )
    """)
    conn.commit()
    conn.close()

init_db()

# ──────────────────────────────────────────────
# CLAHE Preprocessing (matches training pipeline)
# ──────────────────────────────────────────────

def apply_clahe(img_uint8):
    """CLAHE contrast enhancement on LAB L-channel (same as training)."""
    lab = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)
    merged = cv2.merge((cl, a, b))
    return cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)

def preprocess(img_pil):
    """Preprocess a PIL image → model-ready numpy array."""
    img = np.array(img_pil.convert("RGB"))
    img = cv2.resize(img, (300, 300))
    img_clahe = apply_clahe(img.astype(np.uint8))
    from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
    img_final = preprocess_input(img_clahe.astype(np.float32))
    return np.expand_dims(img_final, axis=0)

# ──────────────────────────────────────────────
# Routes — Pages
# ──────────────────────────────────────────────

@app.route("/")
def home():
    return render_template("home.html")

@app.route("/about")
def about():
    return render_template("about.html")

@app.route("/diseases_info")
def diseases_info():
    return render_template("diseases_info.html")

# ──────────────────────────────────────────────
# Routes — Auth
# ──────────────────────────────────────────────

@app.route("/login", methods=["GET", "POST"])
def login():
    if session.get("logged_in"):
        return redirect(url_for("prediction"))

    if request.method == "POST":
        email    = request.form.get("email", "").strip().lower()
        password = request.form.get("password", "")

        conn = get_db()
        user = conn.execute("SELECT * FROM users WHERE email = ?", (email,)).fetchone()
        conn.close()

        if not user:
            flash("No account found with this email. Please sign up first.", "danger")
            return redirect(url_for("login"))

        if not check_password_hash(user["password"], password):
            flash("Incorrect password. Please try again.", "danger")
            return redirect(url_for("login"))

        session.update({
            "logged_in":  True,
            "user_id":    user["id"],
            "user_name":  user["name"],
            "user_email": user["email"]
        })
        flash(f"Welcome back, {user['name']}! 👋", "success")
        return redirect(url_for("prediction"))

    return render_template("login.html")

@app.route("/signup", methods=["POST"])
def signup():
    name     = request.form.get("name", "").strip()
    email    = request.form.get("email", "").strip().lower()
    password = request.form.get("password", "")

    if not name or not email or not password:
        flash("All fields are required.", "danger")
        return redirect(url_for("login"))

    if len(password) < 6:
        flash("Password must be at least 6 characters.", "warning")
        return redirect(url_for("login"))

    conn = get_db()
    existing = conn.execute("SELECT id FROM users WHERE email = ?", (email,)).fetchone()

    if existing:
        conn.close()
        flash("Email already registered. Please log in.", "warning")
        return redirect(url_for("login"))

    conn.execute(
        "INSERT INTO users (name, email, password) VALUES (?, ?, ?)",
        (name, email, generate_password_hash(password))
    )
    conn.commit()
    conn.close()

    flash("Account created! Please log in.", "success")
    return redirect(url_for("login"))

@app.route("/logout")
def logout():
    session.clear()
    flash("Logged out successfully.", "success")
    return redirect(url_for("home"))

# ──────────────────────────────────────────────
# Routes — Prediction
# ──────────────────────────────────────────────

@app.route("/prediction", methods=["GET", "POST"])
def prediction():
    if not session.get("logged_in"):
        flash("Please log in to access predictions.", "warning")
        return redirect(url_for("login"))

    result = None

    if request.method == "POST":
        file = request.files.get("retina_image")

        if not file or file.filename == "":
            flash("Please select an image file.", "warning")
            return redirect(url_for("prediction"))

        # Validate file extension
        allowed = {"jpg", "jpeg", "png", "bmp", "webp"}
        ext = file.filename.rsplit(".", 1)[-1].lower()
        if ext not in allowed:
            flash("Invalid file type. Please upload JPG, PNG or BMP.", "danger")
            return redirect(url_for("prediction"))

        try:
            image = Image.open(file).convert("RGB")
        except Exception:
            flash("Could not read image. Please upload a valid image file.", "danger")
            return redirect(url_for("prediction"))

        # Save uploaded image for preview
        img_filename = f"{uuid.uuid4().hex}.jpg"
        img_save_path = os.path.join(UPLOAD_FOLDER, img_filename)
        image.save(img_save_path, "JPEG", quality=85)
        img_url = url_for("static", filename=f"uploads/{img_filename}")

        # Preprocess & predict
        processed = preprocess(image)
        preds = eye_model.predict(processed, verbose=0)[0]

        # Top prediction
        pred_idx   = int(np.argmax(preds))
        confidence = float(preds[pred_idx]) * 100
        pred_class = CLASS_NAMES.get(pred_idx, "unknown")

        # All 5 class probabilities (sorted desc)
        all_probs = [
            {
                "name":  CLASS_NAMES.get(i, "?").replace("_", " ").title(),
                "key":   CLASS_NAMES.get(i, "?"),
                "prob":  round(float(preds[i]) * 100, 1),
                "color": DISEASE_INFO.get(CLASS_NAMES.get(i, ""), {}).get("color", "#6b7280")
            }
            for i in np.argsort(preds)[::-1]
        ]

        disease_data = DISEASE_INFO.get(pred_class, {
            "display":  pred_class.replace("_", " ").title(),
            "emoji":    "🔍",
            "advice":   "Please consult an eye specialist for proper diagnosis.",
            "severity": "warning",
            "color":    "#f59e0b",
            "tips":     ["Consult an ophthalmologist for detailed examination"]
        })

        result = {
            "disease":    disease_data["display"],
            "emoji":      disease_data["emoji"],
            "confidence": round(confidence, 2),
            "advice":     disease_data["advice"],
            "severity":   disease_data["severity"],
            "color":      disease_data["color"],
            "tips":       disease_data.get("tips", []),
            "all_probs":  all_probs,
            "img_url":    img_url,
        }

    return render_template(
        "prediction.html",
        result=result,
        user_name=session.get("user_name")
    )

# ──────────────────────────────────────────────
# Run
# ──────────────────────────────────────────────

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host="0.0.0.0", port=port, debug=False)