import os
import sqlite3
from collections import Counter
from datetime import datetime
from pathlib import Path
from urllib.parse import urlparse
from urllib.request import Request, urlopen
from uuid import uuid4

import numpy as np
from flask import Flask, flash, redirect, render_template, request, session, url_for
from PIL import Image
from werkzeug.middleware.proxy_fix import ProxyFix
from werkzeug.security import check_password_hash, generate_password_hash
from werkzeug.utils import secure_filename

try:
    from tensorflow.keras.models import load_model
except Exception:
    load_model = None


BASE_DIR = Path(__file__).resolve().parent
UPLOAD_DIR = BASE_DIR / "static" / "uploads"
DATABASE_PATH = BASE_DIR / "instance" / "traffic_signs.sqlite3"
MODEL_PATH = BASE_DIR / "traffic_classifier.h5"
ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "webp"}
NORMALIZE_INPUT = os.environ.get("NORMALIZE_INPUT", "true").lower() in {"1", "true", "yes"}  # le modèle a été entraîné sur des pixels /255
MC_DROPOUT_PASSES = int(os.environ.get("MC_DROPOUT_PASSES", "1"))
UNKNOWN_CONFIDENCE_THRESHOLD = float(os.environ.get("UNKNOWN_CONFIDENCE_THRESHOLD", "0.85"))   # en dessous → "Unknown"
UNKNOWN_MARGIN_THRESHOLD = float(os.environ.get("UNKNOWN_MARGIN_THRESHOLD", "0.10"))             # écart min entre 1er et 2e choix
UNKNOWN_VOTE_THRESHOLD = float(os.environ.get("UNKNOWN_VOTE_THRESHOLD", "0.0"))                  # accord MC-Dropout (0 = désactivé)
URL_CONTENT_TYPES = {
    "image/jpeg": "jpg",
    "image/png": "png",
    "image/webp": "webp",
}

TRAFFIC_SIGN_CLASSES = [
    "Speed limit (20km/h)",
    "Speed limit (30km/h)",
    "Speed limit (50km/h)",
    "Speed limit (60km/h)",
    "Speed limit (70km/h)",
    "Speed limit (80km/h)",
    "End of speed limit (80km/h)",
    "Speed limit (100km/h)",
    "Speed limit (120km/h)",
    "No passing",
    "No passing for vehicles over 3.5 metric tons",
    "Right-of-way at the next intersection",
    "Priority road",
    "Yield",
    "Stop",
    "No vehicles",
    "Vehicles over 3.5 metric tons prohibited",
    "No entry",
    "General caution",
    "Dangerous curve to the left",
    "Dangerous curve to the right",
    "Double curve",
    "Bumpy road",
    "Slippery road",
    "Road narrows on the right",
    "Road work",
    "Traffic signals",
    "Pedestrians",
    "Children crossing",
    "Bicycles crossing",
    "Beware of ice/snow",
    "Wild animals crossing",
    "End of all speed and passing limits",
    "Turn right ahead",
    "Turn left ahead",
    "Ahead only",
    "Go straight or right",
    "Go straight or left",
    "Keep right",
    "Keep left",
    "Roundabout mandatory",
    "End of no passing",
    "End of no passing by vehicles over 3.5 metric tons",
    "UNKNOWN / Rejected"
]


app = Flask(__name__)
app.config["SECRET_KEY"] = os.environ.get("SECRET_KEY", "a7f69cd2-a2a7-4daa-85a0-984f44fcecea")
app.config["MAX_CONTENT_LENGTH"] = 8 * 1024 * 1024
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_prefix=1)

if os.environ.get("SPACE_ID"):
    app.config["SESSION_COOKIE_SECURE"] = True
    app.config["SESSION_COOKIE_SAMESITE"] = "None"
else:
    app.config["SESSION_COOKIE_SAMESITE"] = "Lax"

UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
DATABASE_PATH.parent.mkdir(parents=True, exist_ok=True)

model = None
model_error = None


def get_db():
    conn = sqlite3.connect(DATABASE_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    with get_db() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                email TEXT NOT NULL UNIQUE,
                password_hash TEXT NOT NULL,
                created_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS predictions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                image_path TEXT NOT NULL,
                predicted_class TEXT NOT NULL,
                confidence REAL,
                is_correct INTEGER,
                created_at TEXT NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users(id)
            );
            """
        )


def load_classifier():
    global model, model_error
    if not MODEL_PATH.exists():
        model_error = "traffic_classifier.h5 was not found in the project root."
        return
    if load_model is None:
        model_error = "TensorFlow is not available in this environment."
        return
    try:
        model = load_model(MODEL_PATH)
        model_error = None
    except Exception as exc:
        model_error = f"Model could not be loaded: {exc}"


def current_user():
    user_id = session.get("user_id")
    if not user_id:
        return None
    with get_db() as conn:
        return conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()


def login_required(view):
    def wrapped(*args, **kwargs):
        if not current_user():
            flash("Please log in to access the classifier.", "warning")
            return redirect(url_for("login", next=request.path))
        return view(*args, **kwargs)

    wrapped.__name__ = view.__name__
    return wrapped


def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS


def validate_password_strength(password):
    checks = [
        (len(password) >= 10, "at least 10 characters"),
        (any(char.islower() for char in password), "one lowercase letter"),
        (any(char.isupper() for char in password), "one uppercase letter"),
        (any(char.isdigit() for char in password), "one number"),
        (any(not char.isalnum() for char in password), "one special character"),
    ]
    missing = [message for valid, message in checks if not valid]
    if missing:
        return "Password must contain " + ", ".join(missing) + "."
    return None


def extension_from_url(url, content_type):
    path = urlparse(url).path
    if "." in path:
        extension = path.rsplit(".", 1)[1].lower()
        if extension in ALLOWED_EXTENSIONS:
            return extension
    return URL_CONTENT_TYPES.get((content_type or "").split(";")[0].lower())


def save_remote_image(image_url):
    parsed = urlparse(image_url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        raise ValueError("Please provide a valid HTTP or HTTPS image link.")

    request_data = Request(image_url, headers={"User-Agent": "TrafficSignClassifier/1.0"})
    with urlopen(request_data, timeout=10) as response:
        content_type = response.headers.get("Content-Type", "")
        extension = extension_from_url(image_url, content_type)
        if not extension:
            raise ValueError("The link must point to a PNG, JPG, JPEG, or WEBP image.")

        filename = f"{uuid4().hex}_remote.{extension}"
        save_path = UPLOAD_DIR / filename
        max_bytes = app.config["MAX_CONTENT_LENGTH"]
        total = 0
        with save_path.open("wb") as output:
            while True:
                chunk = response.read(1024 * 64)
                if not chunk:
                    break
                total += len(chunk)
                if total > max_bytes:
                    save_path.unlink(missing_ok=True)
                    raise ValueError("Remote image is too large. Maximum size is 8 MB.")
                output.write(chunk)

    try:
        Image.open(save_path).verify()
    except Exception as exc:
        save_path.unlink(missing_ok=True)
        raise ValueError("The remote file is not a readable image.") from exc

    return save_path, f"uploads/{filename}"


def save_uploaded_image(file):
    original = secure_filename(file.filename)
    filename = f"{uuid4().hex}_{original}"
    save_path = UPLOAD_DIR / filename
    file.save(save_path)
    try:
        Image.open(save_path).verify()
    except Exception as exc:
        save_path.unlink(missing_ok=True)
        raise ValueError("The uploaded file is not a readable image.") from exc
    return save_path, f"uploads/{filename}"


def prepare_image(path):
    image = Image.open(path).convert("RGB")
    image = image.resize((30, 30))
    array = np.asarray(image, dtype=np.float32)
    if NORMALIZE_INPUT:
        array = array / 255.0
    return np.expand_dims(array, axis=0)


def normalize_model_output(predictions):
    predictions = np.asarray(predictions, dtype=np.float64)
    total = float(np.sum(predictions))
    if np.any(predictions < 0) or not np.isclose(total, 1.0, atol=1e-3):
        shifted = predictions - np.max(predictions)
        exp_values = np.exp(shifted)
        return exp_values / np.sum(exp_values)
    return predictions


def predict_probabilities(batch):
    if MC_DROPOUT_PASSES <= 1:
        return normalize_model_output(model.predict(batch, verbose=0)[0]), 1.0

    samples = []
    votes = []
    for _ in range(MC_DROPOUT_PASSES):
        probabilities = normalize_model_output(model(batch, training=True).numpy()[0])
        samples.append(probabilities)
        votes.append(int(np.argmax(probabilities)))

    mean_probabilities = np.mean(np.asarray(samples), axis=0)
    vote_counts = np.bincount(votes, minlength=len(mean_probabilities))
    vote_agreement = float(np.max(vote_counts) / MC_DROPOUT_PASSES)
    return mean_probabilities, vote_agreement


def predict_sign(path):
    if model is None:
        return "Model unavailable", 0.0, False

    probabilities, vote_agreement = predict_probabilities(prepare_image(path))
    ranked_indices   = np.argsort(probabilities)[::-1]
    class_index      = int(ranked_indices[0])
    confidence       = float(probabilities[class_index])
    second_confidence = float(probabilities[int(ranked_indices[1])]) if len(ranked_indices) > 1 else 0.0
    margin           = confidence - second_confidence

    # ── Rejet : image inconnue ou hors-distribution ───────────────────────────
    # Le softmax force toujours une prédiction ; ces trois critères détectent
    # les cas où le modèle n'est pas suffisamment sûr de lui.
    is_unknown = (
        confidence < UNKNOWN_CONFIDENCE_THRESHOLD   # confiance globale trop faible
        or margin < UNKNOWN_MARGIN_THRESHOLD         # trop proche du 2e candidat
        or vote_agreement < UNKNOWN_VOTE_THRESHOLD   # MC-Dropout désaccord
    )

    if is_unknown:
        return "Unknown traffic sign", confidence, True

    label = TRAFFIC_SIGN_CLASSES[class_index] if class_index < len(TRAFFIC_SIGN_CLASSES) else f"Class {class_index}"
    return label, confidence, False


def format_confidence(confidence):
    value = max(0.0, min(float(confidence or 0), 1.0)) * 100
    if 99.995 <= value < 100:
        value = 99.99
    return f"{value:.2f}%"


@app.context_processor
def inject_user():
    return {"format_confidence": format_confidence, "user": current_user()}


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


@app.route("/register", methods=["GET", "POST"])
def register():
    if current_user():
        return redirect(url_for("predict"))

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

        if not name or not email or not password:
            flash("All fields are required.", "danger")
        elif password != confirm_password:
            flash("Passwords do not match.", "danger")
        elif validate_password_strength(password):
            flash(validate_password_strength(password), "danger")
        else:
            try:
                with get_db() as conn:
                    cursor = conn.execute(
                        """
                        INSERT INTO users (name, email, password_hash, created_at)
                        VALUES (?, ?, ?, ?)
                        """,
                        (name, email, generate_password_hash(password), datetime.utcnow().isoformat()),
                    )
                    session["user_id"] = cursor.lastrowid
                    session.permanent = True
                    session.modified = True
                flash("Account created. Welcome to the classifier.", "success")
                return redirect(request.args.get("next") or url_for("predict"))
            except sqlite3.IntegrityError:
                flash("This email is already registered.", "danger")

    return render_template("register.html")


@app.route("/login", methods=["GET", "POST"])
def login():
    if current_user():
        return redirect(url_for("predict"))

    if request.method == "POST":
        email = request.form.get("email", "").strip().lower()
        password = request.form.get("password", "")
        with get_db() as conn:
            user = conn.execute("SELECT * FROM users WHERE email = ?", (email,)).fetchone()
        if user and check_password_hash(user["password_hash"], password):
            session["user_id"] = user["id"]
            session.permanent = True
            session.modified = True
            flash("Connection established.", "success")
            return redirect(request.args.get("next") or url_for("predict"))
        flash("Invalid email or password.", "danger")
    return render_template("login.html")


@app.route("/logout")
def logout():
    session.clear()
    flash("You have been logged out.", "info")
    return redirect(url_for("welcome"))


@app.route("/predict", methods=["GET", "POST"])
@login_required
def predict():
    prediction = None
    if request.method == "POST":
        file = request.files.get("image")
        image_url = request.form.get("image_url", "").strip()
        save_path = None
        image_path = None

        try:
            if file and file.filename:
                if not allowed_file(file.filename):
                    flash("Upload a PNG, JPG, JPEG, or WEBP image.", "danger")
                else:
                    save_path, image_path = save_uploaded_image(file)
            elif image_url:
                save_path, image_path = save_remote_image(image_url)
            else:
                flash("Please choose a traffic sign image or paste an image link.", "warning")
        except ValueError as exc:
            flash(str(exc), "danger")
        except Exception:
            flash("The remote image could not be downloaded. Please try another link.", "danger")

        if save_path and image_path:
            prediction = create_prediction(save_path, image_path)
            if model_error:
                flash(model_error, "warning")

    history = get_user_history(limit=6)
    return render_template("predict.html", prediction=prediction, history=history, model_error=model_error)


def create_prediction(save_path, image_path):
    label, confidence, is_unknown = predict_sign(save_path)
    with get_db() as conn:
        cursor = conn.execute(
            """
            INSERT INTO predictions
            (user_id, image_path, predicted_class, confidence, is_correct, created_at)
            VALUES (?, ?, ?, ?, NULL, ?)
            """,
            (
                session["user_id"],
                image_path,
                label,
                confidence,
                datetime.utcnow().isoformat(),
            ),
        )
        prediction_id = cursor.lastrowid
    return {
        "id": prediction_id,
        "image_path": image_path,
        "predicted_class": label,
        "confidence": confidence,
        "is_unknown": is_unknown,   # utilisé dans le template pour afficher l'avertissement
    }


@app.post("/feedback/<int:prediction_id>")
@login_required
def feedback(prediction_id):
    value = request.form.get("is_correct")
    if value not in {"0", "1"}:
        flash("Feedback must be true or false.", "danger")
        return redirect(url_for("dashboard"))
    with get_db() as conn:
        conn.execute(
            """
            UPDATE predictions
            SET is_correct = ?
            WHERE id = ? AND user_id = ?
            """,
            (int(value), prediction_id, session["user_id"]),
        )
    flash("Feedback saved for the dashboard history.", "success")
    return redirect(request.form.get("next") or url_for("dashboard"))


@app.route("/dashboard")
@login_required
def dashboard():
    history = get_user_history()
    total = len(history)
    reviewed = [row for row in history if row["is_correct"] is not None]
    correct = [row for row in reviewed if row["is_correct"] == 1]
    stats = {
        "total": total,
        "reviewed": len(reviewed),
        "correct": len(correct),
        "incorrect": len(reviewed) - len(correct),
        "accuracy": round((len(correct) / len(reviewed)) * 100, 1) if reviewed else 0,
    }
    charts = build_dashboard_charts(history, stats)
    return render_template("dashboard.html", history=history, stats=stats, charts=charts)


def build_dashboard_charts(history, stats):
    class_counts = Counter(row["predicted_class"] for row in history)
    max_class_count = max(class_counts.values(), default=1)
    class_chart = [
        {
            "label": label,
            "count": count,
            "percent": round((count / max_class_count) * 100, 1),
        }
        for label, count in class_counts.most_common(8)
    ]

    feedback_chart = [
        {"label": "True predictions", "count": stats["correct"], "percent": percent_of(stats["correct"], stats["total"])},
        {"label": "False predictions", "count": stats["incorrect"], "percent": percent_of(stats["incorrect"], stats["total"])},
        {
            "label": "Pending review",
            "count": stats["total"] - stats["reviewed"],
            "percent": percent_of(stats["total"] - stats["reviewed"], stats["total"]),
        },
    ]
    return {"class_chart": class_chart, "feedback_chart": feedback_chart}


def percent_of(value, total):
    if not total:
        return 0
    return round((value / total) * 100, 1)


def get_user_history(limit=None):
    query = """
        SELECT * FROM predictions
        WHERE user_id = ?
        ORDER BY created_at DESC
    """
    params = [session["user_id"]]
    if limit:
        query += " LIMIT ?"
        params.append(limit)
    with get_db() as conn:
        return conn.execute(query, params).fetchall()


init_db()
load_classifier()


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