import os
import random
from ultralytics import YOLO
import cv2
import numpy as np


class PhotoManager:

    def __init__(self):
        self.parent_folder = "/home/luca/git_repos/telegram_amicobot/data/photos"
        self.file_extension = "jpg"

    def get_random_photo(self):
        files = [f for f in os.listdir(self.parent_folder) if os.path.isfile(os.path.join(self.parent_folder, f))]
        if not files:
            raise FileNotFoundError("No files found in the specified folder")

        return os.path.join(self.parent_folder, random.choice(files))
    
    def replace_faces(self, photo_bytes):

        # Convert the byte array to a NumPy array
        nparr = np.frombuffer(photo_bytes, np.uint8)

        # Decode the image
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        original_img = img

        # Specify the desired padding size
        padding_perc = 5
        padding_size = int(img.shape[1] * (1+padding_perc/100))

        # Calculate new dimensions
        new_width = img.shape[1] + 2 * padding_size
        new_height = img.shape[0] + 2 * padding_size

        # Create a blank image with the new dimensions
        padded_image = np.zeros((new_height, new_width, 3), dtype=np.uint8)

        # Copy the original image into the center of the blank image
        padded_image[padding_size:padding_size + img.shape[0], padding_size:padding_size + img.shape[1]] = img
        img = padded_image

        # Load a pre-trained YOLOv8n model
        # Net downloaded from https://github.com/noorkhokhar99/face-detection-yolov8
        model = YOLO('yolov8n-face.pt')
        names = model.model.names

        # Perform inference on 'bus.jpg' with specified parameters with conf=0.5
        results = model.predict(img, verbose=False, conf=0.3, device='cpu')

        # Process detections
        boxes = results[0].boxes.xywh.cpu()
        clss = results[0].boxes.cls.cpu().tolist()
        confs = results[0].boxes.conf.float().cpu().tolist()

        # for box, cls, conf in zip(boxes, clss, confs):
        #     print(f"Class Name: {names[int(cls)]}, Confidence Score: {conf}, Bounding Box: {box}")

        for box, cls, conf in zip(boxes, clss, confs):
            x, y, w, h = map(int, box)
            class_name = names[int(cls)]
            
            # Calculate the bottom-right corner coordinates
            x1, y1 = x - w/2, y - h/2
            x2, y2 = x + w/2, y + h/2

            x1 = int(x1)
            y1 = int(y1)
            x2 = int(x2)
            y2 = int(y2)
            
            # # Draw bounding box
            # cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # # Display class name and confidence
            # label = f"{class_name}: {conf:.2f}"
            # cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            overlay = cv2.imread('/home/luca/git_repos/telegram_amicobot/data/faces/piovra.png', cv2.IMREAD_UNCHANGED)  # Load with alpha channel

            overlay_size_h = overlay.shape[0]
            box_size_h = h
            scale_h = (box_size_h/overlay_size_h) * 100

            overlay_size_w = overlay.shape[1]
            box_size_w = w
            scale_w = (box_size_w/overlay_size_w) * 100
    

            scale_percent = max(scale_h, scale_w) * 1.5 # percent of original size
            width = int(overlay.shape[1] * scale_percent / 100)
            height = int(overlay.shape[0] * scale_percent / 100)
            dim = (width, height)
            
            # resize image
            overlay = cv2.resize(overlay, dim, interpolation = cv2.INTER_AREA)

            h, w = img.shape[:2]
            h1, w1 = overlay.shape[:2]

            # let store center coordinate as cx, cy
            cx, cy = x, y#(w - w1) // 2, (h - h1) // 2  # Note the order of width and height here
            cx = cx - int(overlay.shape[1]/2)
            cy = cy - int(overlay.shape[0]/2) - int(overlay.shape[0] * 0.15)

            if(cx < 0):
                cx = 0
            if(cy < 0):
                cy = 0

            # Ensure the overlay image has an alpha channel
            if overlay.shape[2] == 3:
                overlay = cv2.cvtColor(overlay, cv2.COLOR_BGR2BGRA)

            # Create masks for overlay and background
            overlay_mask = overlay[:, :, 3] / 255.0
            background_mask = 1.0 - overlay_mask

            # Resize overlay image to fit within specified region
            overlay_resized = cv2.resize(overlay, (w1, h1))

            # Blend the images using alpha blending
            for c in range(0, 3):
                img[cy:cy + h1, cx:cx + w1, c] = (overlay_resized[:, :, c] * overlay_mask +
                                                img[cy:cy + h1, cx:cx + w1, c] * background_mask)


        img = img[padding_size:padding_size + original_img.shape[0], padding_size:padding_size + original_img.shape[1]]
        # Convert OpenCV image to bytes
        _, image_bytes = cv2.imencode('.jpg', img)
        image_bytes = image_bytes.tobytes()

        return image_bytes

        # cv2.imshow('Image from Byte Array', img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()