import os
import datetime
from flask import Blueprint, request, jsonify, Response, stream_with_context
from PIL import Image
from ultralytics import YOLO
from config import IMAGE_FOLDER, PRETRAINED_MODEL, TRAINED_MODEL_FOLDER, TEMP_FOLDER
from utils import ensure_correct_permissions, get_latest_model_including_today
import random
import cv2
import json
import time
import subprocess

# Load mô hình YOLO
model = YOLO(get_latest_model_including_today(trained_model_folder=TRAINED_MODEL_FOLDER))

detect_bp = Blueprint('detect', __name__)


connected_clients = {}  # Lưu client_id và kết nối

@detect_bp.route('/detect_images', methods=['POST'])
def detect_images():
    files = request.files.getlist('images')  # Nhận nhiều ảnh từ request
    if not files or len(files) == 0:
        return jsonify({"error": "No files uploaded"}), 400

    results_list = []  # Danh sách kết quả cho từng ảnh


    for file in files:
        try:
            # Kiểm tra file hợp lệ không
            image = Image.open(file)

            # Check kích thước ảnh
            min_width, min_height = 500, 500
            if image.width < min_width or image.height < min_height:
                results_list.append({"filename": file.filename, "error": f"Image too small ({image.width}x{image.height})"})
                continue  # Bỏ qua ảnh lỗi

            # Reset file pointer trước khi lưu
            file.seek(0)

            # Tạo tên file duy nhất
            suffix = datetime.datetime.now().strftime("%y%m%d_%H%M%S")
            image_name = f"detect_{suffix}_{file.filename}"
            image_path = os.path.join("images", image_name)
            file.save(image_path)

            # Chạy mô hình YOLO để phát hiện vật thể
            results = model(image_path, conf=0.6)

            points = []
            for r in results:
                for box in r.boxes:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])  # Lấy tọa độ bbox
                    conf = float(box.conf[0])  # Độ tin cậy
                    cls = int(box.cls[0])  # Nhãn class ID
                    points.append({'label': model.names[cls], 'x1': x1, "y1": y1, "x2": x2, "y2": y2})

            results_list.append({
                "filename": file.filename,
                "image_name": image_name,
                "image_path": image_path,
                "labels": model.names,
                "points": points
            })

        except Exception as e:
            results_list.append({"filename": file.filename, "error": str(e)})

    return jsonify({"results": results_list})


@detect_bp.route('/detect_image', methods=['POST'])
def detect_image():
    file = request.files.get('image')
    if not file:
        return jsonify({"error": "No file uploaded"}), 400

    # Open the image using PIL
    try:
        image = Image.open(file)
    except Exception as e:
        return jsonify({"error": "Invalid image file", "details": str(e)}), 400

    # Check image size
    min_width, min_height = 500, 500
    if image.width < min_width or image.height < min_height:
        return jsonify({"error": f"Image is too small. Minimum size is {min_width}x{min_height} pixels."}), 400

    # Reset file pointer before saving
    file.seek(0)

    # 📌 Get time today
    today_str = datetime.datetime.now().strftime("%Y-%m-%d")

    # 📌 Create folder images/today if not exits
    save_dir = os.path.join(IMAGE_FOLDER, today_str)
    os.makedirs(save_dir, exist_ok=True)
    ensure_correct_permissions(save_dir)
    # 📌 Create radom image name
    basename = "detect_image"
    suffix = datetime.datetime.now().strftime("%H%M%S")  # Chỉ lấy giờ phút giây
    random_number = random.randint(100000, 999999)
    image_name = f"{basename}_{suffix}_{random_number}.png"
    image_path = os.path.join(save_dir, image_name)

    # 📌 Save image to folder today time
    file.save(image_path)

    # Run with model
    results = model(image_path, conf=0.6)  
    points = []

    # Result predict
    for r in results:
        for box in r.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])  
            conf = float(box.conf[0])  
            cls = int(box.cls[0])  

            points.append({'label': model.names[cls], 'x1': x1, "y1": y1, "x2": x2, "y2": y2})

    result = {
        "image_name": image_name,
        "image_path": image_path,
        "labels": model.names,
        "points": points
    }
    return jsonify(result)


@detect_bp.route('/reload_model', methods=['POST'])
def reload_model():
    global model
    new_model_path = get_latest_model_including_today(trained_model_folder=TRAINED_MODEL_FOLDER)
    print(new_model_path)
    model = YOLO(new_model_path)  
    return jsonify({"message": f"Model reloaded from {new_model_path}", "status": "true"}), 200

@detect_bp.route('/train', methods=['GET'])
def run_training():
    try:
        # Run train.py as a subprocess
        process = subprocess.Popen(["python", "train.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()

        if process.returncode == 0:
            return jsonify({"message": "Training started successfully", "output": stdout.decode()}), 200
        else:
            return jsonify({"error": "Training failed", "details": stderr.decode()}), 500

    except Exception as e:
        return jsonify({"error": str(e)}), 500


