from flask import Flask, render_template, request, jsonify, Response
from ultralytics import YOLO
import cv2
import numpy as np
import base64
import os
from PIL import Image
import io
import threading

# ── TOGGLE: set USE_VLM = True to enable VLM analysis ────────────────────────
USE_VLM = True
# ─────────────────────────────────────────────────────────────────────────────

if USE_VLM:
    from vlm_analyzer import MicroplasticRiskAnalyzer

app = Flask(__name__)

# ── YOLO MODEL PATH ───────────────────────────────────────────────────────────
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "best.pt")
# ─────────────────────────────────────────────────────────────────────────────

CONFIDENCE  = 0.25
IOU_THRESH  = 0.45
IMAGE_SIZE  = 640

# Load YOLO once at startup
print("Loading YOLO model...")
model       = YOLO(MODEL_PATH)
CLASS_NAMES = model.names
print(f"YOLO loaded. Classes: {CLASS_NAMES}")

# Load VLM if enabled
vlm = None
if USE_VLM:
    print("Loading VLM (Qwen2-VL-2B-Instruct)...")
    vlm = MicroplasticRiskAnalyzer()
    print("VLM loaded and ready.")

# Webcam state
webcam_active = False
webcam_cap    = None
webcam_lock   = threading.Lock()


# ── HELPERS ───────────────────────────────────────────────────────────────────

def run_detection(image_array):
    """Run YOLO detection. Returns annotated image + detections list."""
    results = model.predict(
        source  = image_array,
        conf    = CONFIDENCE,
        iou     = IOU_THRESH,
        imgsz   = IMAGE_SIZE,
        verbose = False,
    )[0]

    annotated  = results.plot(line_width=2)
    detections = []

    if results.boxes is not None and len(results.boxes):
        for box in results.boxes:
            detections.append({
                "class":      CLASS_NAMES[int(box.cls[0])],
                "confidence": round(float(box.conf[0]) * 100, 1),
                "bbox":       [round(v) for v in box.xyxy[0].tolist()]
            })

    return annotated, detections


def numpy_to_base64(img_array):
    """Convert BGR numpy image → base64 JPEG string."""
    img_rgb = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(img_rgb)
    buffer  = io.BytesIO()
    pil_img.save(buffer, format="JPEG", quality=90)
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def run_vlm_analysis(image_array, detections):
    """Run VLM scene analysis if enabled. Returns dict or None."""
    if not USE_VLM or vlm is None:
        return None
    try:
        return vlm.analyze_scene(image_array, detections)
    except Exception as e:
        print(f"[VLM ERROR] {e}")
        return {
            "risk_level":      "Error",
            "explanation":     str(e),
            "recommendations": "VLM analysis failed. Check logs.",
            "raw":             "",
        }


# ── ROUTES ────────────────────────────────────────────────────────────────────

@app.route("/")
def index():
    return render_template(
        "index.html",
        class_names = list(CLASS_NAMES.values()),
        vlm_enabled = USE_VLM,
    )


@app.route("/detect/image", methods=["POST"])
def detect_image():
    """YOLO detection (+ optional VLM) on uploaded image."""
    if "file" not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file  = request.files["file"]
    data  = np.frombuffer(file.read(), np.uint8)
    image = cv2.imdecode(data, cv2.IMREAD_COLOR)

    if image is None:
        return jsonify({"error": "Could not read image"}), 400

    annotated, detections = run_detection(image)
    vlm_result            = run_vlm_analysis(image, detections)

    response = {
        "image":      numpy_to_base64(annotated),
        "detections": detections,
        "count":      len(detections),
        "vlm":        vlm_result,          # None when VLM disabled
    }
    return jsonify(response)


@app.route("/detect/video_frame", methods=["POST"])
def detect_video_frame():
    """YOLO detection (+ optional VLM) on a single video frame sent as base64."""
    data = request.get_json()
    if not data or "frame" not in data:
        return jsonify({"error": "No frame data"}), 400

    # Strip data-URL prefix if present
    raw_b64    = data["frame"].split(",")[-1]
    frame_data = base64.b64decode(raw_b64)
    nparr      = np.frombuffer(frame_data, np.uint8)
    image      = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if image is None:
        return jsonify({"error": "Could not decode frame"}), 400

    annotated, detections = run_detection(image)
    vlm_result            = run_vlm_analysis(image, detections)

    return jsonify({
        "image":      numpy_to_base64(annotated),
        "detections": detections,
        "count":      len(detections),
        "vlm":        vlm_result,
    })


@app.route("/classes")
def get_classes():
    return jsonify({
        "classes":     list(CLASS_NAMES.values()),
        "total":       len(CLASS_NAMES),
        "vlm_enabled": USE_VLM,
    })


# ── WEBCAM STREAM (server-side, optional) ─────────────────────────────────────

@app.route("/webcam/start", methods=["POST"])
def webcam_start():
    global webcam_active, webcam_cap
    with webcam_lock:
        if webcam_active:
            return jsonify({"status": "already_running"})
        webcam_cap = cv2.VideoCapture(0)
        if not webcam_cap.isOpened():
            webcam_cap = None
            return jsonify({"error": "Cannot open webcam"}), 500
        webcam_active = True
    return jsonify({"status": "started"})


@app.route("/webcam/stop", methods=["POST"])
def webcam_stop():
    global webcam_active, webcam_cap
    with webcam_lock:
        webcam_active = False
        if webcam_cap is not None:
            webcam_cap.release()
            webcam_cap = None
    return jsonify({"status": "stopped"})


# ── ENTRY POINT ───────────────────────────────────────────────────────────────

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