routes/classification_routes.py

import os, io, base64, requests, traceback
from pathlib import Path
from flask import Blueprint, request, jsonify
import numpy as np
from PIL import Image
import tensorflow as tf

# =============================================================
# ✅ Import MobilenetV2 utilities
# =============================================================
try:
    from tensorflow.keras.applications import mobilenet_v2 as _mv2
except Exception:
    from keras.applications import mobilenet_v2 as _mv2

preprocess_input = _mv2.preprocess_input
classification_bp = Blueprint("classification_bp", __name__)

# =============================================================
# 🧠 Model Setup (auto-download for Hugging Face)
# =============================================================
MODEL_DIR = Path(os.getenv("MODEL_DIR", "/tmp/model"))
MODEL_DIR.mkdir(parents=True, exist_ok=True)

MODEL_URLS = {
    "model": "https://huggingface.co/Gwen01/ProctorVision-Models/resolve/main/cheating_mobilenetv2_tf15.h5",
    "threshold": "https://huggingface.co/Gwen01/ProctorVision-Models/resolve/main/best_threshold.npy",
}

# --- Download model & threshold if not cached ---
for key, url in MODEL_URLS.items():
    local_path = MODEL_DIR / Path(url).name
    if not local_path.exists():
        print(f"📥 Downloading {key} from {url}")
        try:
            r = requests.get(url, timeout=60)
            r.raise_for_status()
            with open(local_path, "wb") as f:
                f.write(r.content)
            print(f"✅ Saved {key} → {local_path}")
        except Exception as e:
            print(f"❌ Failed to download {key}: {e}")
    else:
        print(f"🗂️ Using cached {key} → {local_path}")

# =============================================================
# 🧩 Model Loading
# =============================================================
model = None
CANDIDATES = [
    "cheating_mobilenetv2_tf15.h5",
    "cheating_mobilenetv2_final.h5",
    "cheating_mobilenetv2_legacy.h5",
    "cheating_mobilenetv2_final.keras",
]

try:
    for name in CANDIDATES:
        path = MODEL_DIR / name
        if not path.exists():
            continue
        try:
            model = tf.keras.models.load_model(path, compile=False)
            print(f"✅ Model loaded successfully from {path}")
            break
        except Exception as e:
            print(f"❌ Could not load {path}: {e}")
            traceback.print_exc()
    if model is None:
        print("🚨 Model could not be initialized — classification routes will return 500 errors.")
except Exception as e:
    print(f"❌ Error during model loading: {e}")
    traceback.print_exc()

# --- Print available files ---
print("📂 MODEL_DIR contents:")
for p in MODEL_DIR.glob("*"):
    print(f"   - {p.name} ({p.stat().st_size} bytes)")

# =============================================================
# 🔹 Load threshold (with fallback)
# =============================================================
thr_file = MODEL_DIR / "best_threshold.npy"
THRESHOLD = 0.555  # default fallback
try:
    if thr_file.exists():
        THRESHOLD = float(np.load(thr_file)[0])
        print(f"✅ Loaded threshold from {thr_file}: {THRESHOLD}")
    else:
        print("⚠️ Threshold file missing — using default 0.555")
except Exception as e:
    print(f"⚠️ Failed to load threshold: {e} — using default 0.555")

# =============================================================
# 📏 Input Shape
# =============================================================
if model is not None:
    H, W = model.input_shape[1:3]
else:
    H, W = 224, 224

LABELS = ["Cheating", "Not Cheating"]

# =============================================================
# 🧩 Helper Functions
# =============================================================
def preprocess_pil(pil_img: Image.Image) -> np.ndarray:
    """Convert PIL image to model-ready numpy array."""
    img = pil_img.convert("RGB").resize((W, H))
    x = np.asarray(img, dtype=np.float32)
    x = preprocess_input(x)
    return np.expand_dims(x, 0)

def predict_batch(batch_np: np.ndarray) -> np.ndarray:
    """Predict probabilities for a batch of images."""
    try:
        raw = model.predict(batch_np, verbose=0)
        if raw.ndim == 2 and raw.shape[1] == 2:
            return raw[:, 1]  # probability of “Not Cheating”
        return raw.ravel()
    except Exception as e:
        print("💥 Prediction error:", e)
        traceback.print_exc()
        raise

def label_from_prob(prob_non_cheating: float) -> str:
    """Map probability to human-readable label."""
    return LABELS[int(prob_non_cheating >= THRESHOLD)]

# =============================================================
# 🔹 Route 1 – Manual Upload (Testing)
# =============================================================
@classification_bp.route("/classify_multiple", methods=["POST"])
def classify_multiple():
    """Classify uploaded image(s)."""
    try:
        if model is None:
            return jsonify({"error": "Model not loaded"}), 500

        files = request.files.getlist("files")
        if not files:
            return jsonify({"error": "No files uploaded"}), 400

        batch = [preprocess_pil(Image.open(io.BytesIO(f.read())))[0] for f in files]
        batch_np = np.stack(batch)
        probs = predict_batch(batch_np)
        labels = [label_from_prob(p) for p in probs]

        return jsonify({
            "threshold": THRESHOLD,
            "results": [
                {"label": lbl, "prob_non_cheating": float(p)}
                for lbl, p in zip(labels, probs)
            ]
        })
    except Exception as e:
        print("💥 classify_multiple crashed:", e)
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500

# =============================================================
# 🔹 Route 2 – Backend-to-Backend Integration (Railway)
# =============================================================
@classification_bp.route("/classify_behavior_logs", methods=["POST"])
def classify_behavior_logs():
    """Fetch behavior logs from Railway backend, classify them, and send results back."""
    try:
        if model is None:
            print("❌ Model not loaded.")
            return jsonify({"error": "Model not loaded"}), 500

        data = request.get_json(silent=True) or {}
        user_id, exam_id = data.get("user_id"), data.get("exam_id")
        if not user_id or not exam_id:
            return jsonify({"error": "Missing user_id or exam_id"}), 400

        RAILWAY_API = os.getenv("RAILWAY_API", "").rstrip("/")
        if not RAILWAY_API:
            print("❌ RAILWAY_API not configured.")
            return jsonify({"error": "RAILWAY_API not configured"}), 500

        # --- Fetch logs ---
        fetch_url = f"{RAILWAY_API}/api/fetch_behavior_logs"
        print(f"📡 Fetching logs from {fetch_url}?user_id={user_id}&exam_id={exam_id}")
        res = requests.get(fetch_url, params={"user_id": user_id, "exam_id": exam_id}, timeout=30)
        print(f"📨 Railway response: {res.status_code}")
        res.raise_for_status()

        logs = res.json().get("logs", [])
        if not logs:
            print("ℹ️ No logs found for classification.")
            return jsonify({"message": "No logs to classify."}), 200

        # --- Classify each log ---
        updates = []
        for log in logs:
            try:
                img_data = base64.b64decode(log["image_base64"])
                pil = Image.open(io.BytesIO(img_data))
                prob = predict_batch(preprocess_pil(pil))[0]
                lbl = label_from_prob(prob)
                updates.append({"id": log["id"], "label": lbl})
            except Exception as e:
                print(f"⚠️ Skipped log {log.get('id')}: {e}")
                traceback.print_exc()

        # --- Send results back to Railway ---
        update_url = f"{RAILWAY_API}/api/update_classifications"
        print(f"📤 Sending {len(updates)} updates to {update_url}")
        post_res = requests.post(update_url, json={"updates": updates}, timeout=30)
        print("📨 Update response:", post_res.status_code, post_res.text[:200])
        post_res.raise_for_status()

        return jsonify({
            "message": f"Classification complete for {len(updates)} logs.",
            "threshold": THRESHOLD
        }), 200

    except Exception as e:
        print("💥 classify_behavior_logs crashed:", e)
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500