from collections import Counter
from datetime import datetime
import os
from PIL import Image
from ultralytics import YOLO
import io

face_model_path = "model/face_detection.pt"
visibility_model_path = "model/visibility.pt"
availability_model_path = "model/availability.pt"

ANNOTATED_FOLDER_VISIBILITY = "uploads/annotated/visibility"
ANNOTATED_FOLDER_AVAILABILITY = "uploads/annotated/availability"
ANNOTATED_FOLDER_OWNER_OUTLET = "uploads/owner_outlet"
os.makedirs(ANNOTATED_FOLDER_VISIBILITY, exist_ok=True)
os.makedirs(ANNOTATED_FOLDER_AVAILABILITY, exist_ok=True)
os.makedirs(ANNOTATED_FOLDER_OWNER_OUTLET, exist_ok=True)

try:
    model_availability = YOLO(availability_model_path)
    model_face = YOLO(face_model_path) 
    model_visibility = YOLO(visibility_model_path)
except Exception as e:
    raise RuntimeError(f"Error loading model: {e}")

class InferenceService:
    @staticmethod
    async def perform_visibility_inference(image_data: bytes, temp_path_visibility: str) -> dict:
        try:
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
            temp_path = temp_path_visibility
            image.save(temp_path, "JPEG", quality=100)

            results = model_visibility(temp_path)
            
            if not results or not results[0].boxes:
                os.remove(temp_path)
                return {"success": False, "error": ["Gambar visibility tidak terdeteksi."]}
            
            confidence_threshold = 0.4  
            boxes = results[0].boxes  
            filtered_boxes = [box for box in boxes if box.conf > confidence_threshold]  

            if not filtered_boxes:
                os.remove(temp_path)
                return {"success": False, "error": ["Gambar visibility tidak terdeteksi."]}

            label_counts = Counter()
            for box in filtered_boxes:
                class_label = int(box.cls)
                if class_label not in model_visibility.names:
                    os.remove(temp_path)
                    return {"success": False, "error": [f"Unknown class label: {class_label}"]}
                label_name = model_visibility.names[class_label]
                label_counts[label_name] += 1

            results[0].boxes = filtered_boxes
            pred_img = results[0].plot()  

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            annotated_path = os.path.join(ANNOTATED_FOLDER_VISIBILITY, f"annotated_{timestamp}.jpg")
            Image.fromarray(pred_img[..., ::-1]).save(annotated_path)

            os.remove(temp_path)

            return {
                "success": True,
                "label_counts": dict(label_counts),
                "annotated_image_path": annotated_path
            }

        except Exception as e:
            if os.path.exists(temp_path_visibility):
                os.remove(temp_path_visibility)
            return {"success": False, "error": [f"Error during inference: {str(e)}"]}

    @staticmethod
    async def perform_availability_inference(image_data: bytes, temp_path_availability: str) -> dict:
        try:
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
            temp_path = temp_path_availability
            image.save(temp_path, "JPEG", quality=100)

            results = model_availability(temp_path)
            
            if not results or not results[0].boxes:
                os.remove(temp_path)
                return {"success": False, "error": ["Gambar availability tidak terdeteksi."]}

            confidence_threshold = 0.4  
            boxes = results[0].boxes  
            filtered_boxes = [box for box in boxes if box.conf > confidence_threshold]  

            if not filtered_boxes:
                os.remove(temp_path)
                return {"success": False, "error": ["Gambar availability tidak terdeteksi."]}

            label_counts = Counter()
            for box in filtered_boxes:
                class_label = int(box.cls)
                if class_label not in model_availability.names:
                    os.remove(temp_path)
                    return {"success": False, "error": [f"Unknown class label: {class_label}"]}
                label_name = model_availability.names[class_label]
                label_counts[label_name] += 1

            results[0].boxes = filtered_boxes
            pred_img = results[0].plot()  

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            annotated_path = os.path.join(ANNOTATED_FOLDER_AVAILABILITY, f"annotated_{timestamp}.jpg")
            Image.fromarray(pred_img[..., ::-1]).save(annotated_path)

            os.remove(temp_path)

            return {
                "success": True,
                "label_counts": dict(label_counts),
                "annotated_image_path": annotated_path
            }

        except Exception as e:
            if os.path.exists(temp_path_availability):
                os.remove(temp_path_availability)
            return {"success": False, "error": [f"Error during inference: {str(e)}"]}

    @staticmethod
    async def perform_face_inference(image_data: bytes) -> dict:
        try:
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
        except Exception as e:
            return {"success": False, "error": [f"Error processing image: {str(e)}"]}

        try:
            results = model_face(image)

            if not results or not results[0].boxes:
                return {"success": False, "error": ["muka tidak terdeteksi."]}

            detections = results[0]

        except Exception as e:
            return {"success": False, "error": [f"Error during inference: {str(e)}"]}

        confidence_threshold = 0.5
        filtered_boxes = [box for box in detections.boxes if box.conf > confidence_threshold]

        if not filtered_boxes:
            return {"success": False, "error": ["muka tidak terdeteksi."]}

        if len(filtered_boxes) < 2:
            return {"success": False, "error": ["Harus ada minimal 2 muka yang terdeteksi."]}

        label_counts = Counter()
        for box in filtered_boxes:
            class_label = int(box.cls)
            if class_label not in model_face.names:
                return {"success": False, "error": [f"Unknown class label: {class_label}"]}
            label_name = model_face.names[class_label]
            label_counts[label_name] += 1

        try:
            detections.boxes = filtered_boxes  
            annotated_array = detections.plot()
            if annotated_array.shape[-1] != 3:
                raise ValueError("Annotated array does not have 3 channels (RGB format expected).")
            annotated_image = Image.fromarray(annotated_array[..., ::-1])

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            annotated_path = os.path.join(ANNOTATED_FOLDER_OWNER_OUTLET, f"annotated_{timestamp}.jpg")
            annotated_image.save(annotated_path)
        except Exception as e:
            return {"success": False, "error": [f"Error saving annotated image: {str(e)}"]}

        return {
            "success": True,
            "label_counts": dict(label_counts),
            "annotated_image_path": annotated_path,
        }