"""
Flask Web Application for YOLOv12 Face Detection (ONNX Backend)
Supports image upload, video upload, and live webcam streaming
Optimized for Headless CPU Cloud environments (Hugging Face Spaces)
"""

import os

os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("ONNXRUNTIME_INTER_OP_NUM_THREADS", "1")
os.environ.setdefault("ONNXRUNTIME_INTRA_OP_NUM_THREADS", "2")

import base64
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, cast

import cv2
import numpy as np
from dotenv import load_dotenv
from flask import Flask, jsonify, render_template, request, send_file
from flask_limiter import Limiter
from flask_limiter.errors import RateLimitExceeded
from flask_limiter.util import get_remote_address
from flask_sqlalchemy import SQLAlchemy
from werkzeug.utils import secure_filename

try:
    from turbojpeg import TurboJPEG

    jpeg_engine = TurboJPEG()
except Exception:
    jpeg_engine = None


def decode_image_fast(image_bytes: bytes) -> np.ndarray:
    """Decode images using TurboJPEG if have, else OpenCV"""
    if jpeg_engine is not None:
        try:
            return jpeg_engine.decode(image_bytes)
        except Exception:
            pass
    nparr = np.frombuffer(image_bytes, np.uint8)
    return cv2.imdecode(nparr, cv2.IMREAD_COLOR)


def encode_jpeg_fast(image_bgr: np.ndarray, quality: int = 85) -> bytes:
    """Zip to TurboJPEG"""
    if jpeg_engine is not None:
        try:
            res = jpeg_engine.encode(image_bgr, quality=quality)
            if res is not None:
                return cast(bytes, res)
        except Exception:
            pass
    _, buffer = cv2.imencode(".jpg", image_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
    return bytes(buffer)


from face_detection_yolov12 import YOLOv12FaceDetector, detect_from_video

load_dotenv()

app = Flask(__name__, template_folder="../web/templates")
app.config["SQLALCHEMY_DATABASE_URI"] = os.getenv("DB_URL")
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "fallback_secret_key_neu_khong_co")
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {"pool_recycle": 280}

db = SQLAlchemy(app)
PROJECT_ROOT = Path(__file__).parent.parent
UPLOAD_FOLDER = PROJECT_ROOT / "data" / "uploads"
MODELS_DIR = PROJECT_ROOT / "models"
ALLOWED_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "mp4", "avi", "mov", "mkv"}
MAX_FILE_SIZE = 1024 * 1024 * 1024  # 1GB

ALLOWED_MODELS = {
    "yolov12n-face-int8.onnx",
    "yolov12s-face-int8.onnx",
    "yolov12m-face-int8.onnx",
    "yolov12l-face-int8.onnx",
}
UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)

app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config["MAX_CONTENT_LENGTH"] = MAX_FILE_SIZE

limiter = Limiter(
    key_func=get_remote_address,
    app=app,
    default_limits=["10000 per day", "1000 per hour"],
    storage_uri="memory://",
)


class Feedback(db.Model):  # type: ignore
    __tablename__ = "feedbacks"
    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(255))
    model_name = db.Column(db.String(50))
    rating = db.Column(db.Integer)
    comment = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)


with app.app_context():
    try:
        db.create_all()
        print("Connect to MySQL successfully!")
    except Exception as e:
        print(f"Connection error: {e}")


def get_detector(model_name):
    """
    Get or create detector instance.
    The YOLOv12FaceDetector class natively handles Singleton caching internally,
    so we don't need a manual dictionary cache here anymore.
    """
    safe_name = secure_filename(model_name)

    if safe_name.endswith(".pt"):
        safe_name = safe_name.replace(".pt", ".onnx")

    if safe_name not in ALLOWED_MODELS:
        logging.error(f"Attempt to load unsupported model: {safe_name}")
        raise ValueError(f"Unsupported model: {safe_name}")

    model_path = MODELS_DIR / safe_name

    try:
        final_path = model_path.resolve()
        safe_root = MODELS_DIR.resolve()
        if not str(final_path).startswith(str(safe_root)):
            logging.error(f"Security Alert: Symlink attack detected! {final_path}")
            raise ValueError("Invalid model path (Symlink violation)")

    except Exception as e:
        logging.error(f"Error resolving model path: {str(e)}")
        raise FileNotFoundError(f"Model path error: {str(e)}")

    if not final_path.exists():
        logging.error(f"Model file not found: {final_path}")
        raise FileNotFoundError(f"Model not found: {final_path}")

    return YOLOv12FaceDetector(final_path.name)


def allowed_file(filename):
    """Check if file extension is allowed"""
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS


def is_image(filename):
    """Check if file is image"""
    ext = filename.rsplit(".", 1)[1].lower()
    return ext in {"jpg", "jpeg", "png", "gif"}


def is_video(filename):
    """Check if file is video"""
    ext = filename.rsplit(".", 1)[1].lower()
    return ext in {"mp4", "avi", "mov", "mkv"}


@app.route("/")
def index():
    """Main page"""
    return render_template("index.html")


@app.route("/api/detect-image", methods=["POST"])
def detect_image():
    """Detect faces in uploaded image or live webcam frame"""
    try:
        if "file" not in request.files:
            return jsonify({"error": "No file provided"}), 400

        file = request.files["file"]
        if file.filename == "":
            return jsonify({"error": "No file selected"}), 400

        if not allowed_file(file.filename) or not is_image(file.filename):
            return jsonify({"error": "Only image files allowed"}), 400

        # Get parameters
        model = request.form.get("model", "yolov12s-face-int8.onnx")
        blur = request.form.get("blur") == "true"

        # Nhận diện xem request có phải đến từ luồng Webcam hay không
        is_webcam = request.form.get("mode") == "webcam" or "webcam" in file.filename.lower()

        if not model.endswith(".onnx"):
            model = model.replace(".pt", ".onnx")

        if model not in ALLOWED_MODELS:
            app.logger.info(f"Invalid model '{model}' requested. Fallback to default.")
            model = "yolov12n-face-int8.onnx" if is_webcam else "yolov12s-face-int8.onnx"

        # Lấy instance detector (đã được cache ở RAM qua Singleton)
        detector = get_detector(model)

        # Giải mã ảnh trực tiếp từ bytes tốc độ cao
        image_data = file.read()
        image = decode_image_fast(image_data)
        if image is None:
            return jsonify({"error": "Invalid image data"}), 400

        img_h, img_w = image.shape[:2]
        detections = detector.detect_faces(image, conf_threshold=0.32)

        # --- NHÁNH 1: WEBCAM (CHỈ TRẢ VỀ TỌA ĐỘ - ZERO SERVER-SIDE RENDERING) ---
        if is_webcam:
            return jsonify(
                {
                    "success": True,
                    "mode": "coordinates_only",
                    "frame_width": img_w,
                    "frame_height": img_h,
                    "detections": {
                        "count": len(detections),
                        "boxes": [
                            {
                                "x1": det["x1"],
                                "y1": det["y1"],
                                "x2": det["x2"],
                                "y2": det["y2"],
                                "w": det["w"],
                                "h": det["h"],
                                "conf": round(det["confidence"], 3),
                            }
                            for det in detections
                        ],
                    },
                }
            )

        # --- NHÁNH 2: ẢNH TĨNH UPLOAD THÔNG THƯỜNG ---
        if blur:
            result_image = detector.draw_faces(image, detections, show_conf=False, blur_faces=blur)
        else:
            result_image = detector.draw_faces(image, detections, show_conf=True, blur_faces=blur)

        if result_image is None:
            return jsonify({"error": "Failed to process image"}), 500

        # Extract crops for gallery
        crops_base64 = []
        if len(detections) > 0:
            crops = detector.get_face_crops(image, detections)
            for crop in crops:
                crop_bytes = encode_jpeg_fast(crop, quality=80)
                crops_base64.append(base64.b64encode(crop_bytes).decode())

        # Convert result to base64 for display
        result_bytes = encode_jpeg_fast(result_image, quality=85)
        img_base64 = base64.b64encode(result_bytes).decode()

        # Prepare response
        response = {
            "success": True,
            "image": f"data:image/jpeg;base64,{img_base64}",
            "detections": {
                "count": len(detections),
                "faces": [
                    {
                        "id": i + 1,
                        "confidence": f"{det['confidence']:.2%}",
                        "width": det["w"],
                        "height": det["h"],
                        "position": f"({det['x1']}, {det['y1']})",
                    }
                    for i, det in enumerate(detections)
                ],
                "crops": crops_base64,
            },
        }

        return jsonify(response)

    except Exception:
        logging.exception("Error during image detection")
        return jsonify({"error": "Internal server error during image detection"}), 500


@app.route("/api/detect-video", methods=["POST"])
@limiter.limit("5 per hour")
def detect_video():
    """Detect faces in uploaded video"""
    try:
        if "file" not in request.files:
            return jsonify({"error": "No file provided"}), 400

        file = request.files["file"]
        if file.filename == "":
            return jsonify({"error": "No file selected"}), 400

        if not allowed_file(file.filename) or not is_video(file.filename):
            return jsonify({"error": "Only video files allowed"}), 400

        # Get model selection
        raw_model = request.form.get("model", "yolov12n-face-int8.onnx")
        model = Path(raw_model).name
        if model.endswith(".pt"):
            model = model.replace(".pt", ".onnx")
        if not model.endswith(".onnx"):
            model = f"{model}.onnx"

        if model not in ALLOWED_MODELS:
            app.logger.info(f"Invalid model '{model}' requested. Fallback to default (Nano INT8).")
            model = "yolov12n-face-int8.onnx"

        # Save uploaded file
        filename = secure_filename(file.filename)
        input_path = UPLOAD_FOLDER / f"input_{filename}"
        output_path = UPLOAD_FOLDER / f"output_{filename}"
        file.save(input_path)

        # Detect faces in video
        detect_from_video(
            video_path=str(input_path),
            model_path=model,
            output_path=str(output_path),
            conf_threshold=0.32,
        )

        # Return file info
        response = {
            "success": True,
            "message": "Video processing complete",
            "output_file": output_path.name,
            "download_url": f"/api/download/{output_path.name}",
        }

        return jsonify(response)

    except Exception:
        app.logger.exception("Error while processing video detection request")
        return jsonify({"error": "Internal server error"}), 500


@app.route("/api/download/<filename>", methods=["GET"])
def download_file(filename):
    try:
        filepath = UPLOAD_FOLDER / secure_filename(filename)
        if not filepath.exists():
            return jsonify({"error": "File not found"}), 404

        as_attachment = request.args.get("download") == "1"

        return send_file(
            filepath,
            mimetype="video/mp4",
            as_attachment=as_attachment,
            conditional=True,
        )
    except Exception:
        app.logger.exception("Error while processing download request for %s", filename)
        return jsonify({"error": "Internal server error"}), 500


@app.route("/api/models", methods=["GET"])
def get_models():
    """Get ALL available ONNX models for dropdown selection"""
    models = {
        "nano": {
            "name": "yolov12n-face-int8.onnx",
            "label": "Nano (n) - Fastest (ONNX INT8)",
            "description": "Real-time speed, best for CPU/Webcam",
            "size": "Smallest",
        },
        "small": {
            "name": "yolov12s-face-int8.onnx",
            "label": "Small (s) - Balanced (ONNX INT8)",
            "description": "Good balance of speed and accuracy",
            "size": "Small",
        },
        "medium": {
            "name": "yolov12m-face-int8.onnx",
            "label": "Medium (m) - High Precision (ONNX INT8)",
            "description": "High accuracy, optimized CPU execution",
            "size": "Medium",
        },
        "large": {
            "name": "yolov12l-face-int8.onnx",
            "label": "Large (l) - Max Accuracy (ONNX INT8)",
            "description": "Best detection quality",
            "size": "Large",
        },
    }

    available: dict[str, Any] = {}
    for key, info in models.items():
        model_path = MODELS_DIR / info["name"]
        if model_path.exists():
            available[key] = info

    order = ["nano", "small", "medium", "large"]
    sorted_available = {k: available[k] for k in order if k in available}

    return jsonify(sorted_available)


@app.route("/api/feedback", methods=["POST"])
@limiter.limit("5 per minute")
def submit_feedback():
    try:
        data = request.json

        new_fb = Feedback(
            filename=data.get("filename"),
            model_name=data.get("model"),
            rating=data.get("rating"),
            comment=data.get("comment", ""),
        )
        db.session.add(new_fb)
        db.session.commit()
        return jsonify({"success": True, "message": "Rating saved!"})
    except Exception as e:
        app.logger.error(f"DB Error: {e}")
        return jsonify({"error": "Database error"}), 500


@app.route("/api/health", methods=["GET"])
def health_check():
    """Health check endpoint"""
    return jsonify({"status": "healthy", "service": "Face Detection ONNX API"})


@app.errorhandler(413)
def request_entity_too_large(error):
    """Handle file size exceeded"""
    return jsonify({"error": "File too large. Maximum 500MB allowed"}), 413


@app.errorhandler(500)
def internal_error(error):
    """Handle internal server error"""
    return jsonify({"error": "Internal server error"}), 500


@app.errorhandler(RateLimitExceeded)
def handle_rate_limit_error(e):
    """Handle rate limit exceeded errors"""
    app.logger.warning(f"Rate limit exceeded: {e.description}")
    return (
        jsonify(
            {
                "error": "Too many requests",
                "message": f"Too fast! Please wait a moment. ({e.description})",
            }
        ),
        429,
    )


if __name__ == "__main__":
    print("\n" + "=" * 70)
    print("🌐 Starting YOLOv12 Face Detection Web Server (ONNX Optimized)")
    print("=" * 70)
    print("\n📍 Server: http://localhost:7860")
    print("📁 Upload folder: ", UPLOAD_FOLDER)
    print("🔧 Models folder: ", MODELS_DIR)
    print("\n🎯 Available endpoints:")
    print("   GET  /                    - Web interface")
    print("   POST /api/detect-image    - Detect faces in image")
    print("   POST /api/detect-video    - Detect faces in video")
    print("   GET  /api/models          - Get available ONNX models")
    print("   GET  /api/health          - Health check")
    print("\n" + "=" * 70 + "\n")

    # Determine debug mode from environment (default: disabled)
    debug_mode = os.getenv("FLASK_ENV") == "development"
    app.run(host="0.0.0.0", port=7860, debug=debug_mode, use_reloader=False)