import folder_paths
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from ultralytics import YOLO
import torch
import os
from urllib.parse import urlparse
import logging
from torch.hub import download_url_to_file
import cv2
import json

logger = logging.getLogger("Comfyui-Yolov8-JSON")
yolov8_model_dir_name = "yolov8"
yolov8_model_list = {
    "yolov8n(6.23MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8n.pt",
    },
    "yolov8s(21.53MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s.pt",
    },
    "yolov8m (49.70MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8m.pt",
    },
    "yolov8l (83.70MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8l.pt",
    },
    "yolov8x (130.53)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8x.pt",
    },
    "yolov8n-seg (6.73MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8n-seg.pt",
    },
    "yolov8s-seg(22.79MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s-seg.pt",
    },
    "yolov8m-seg  (52.36MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8m-seg.pt",
    },
    "yolov8l-seg  (88.11MB)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8l-seg.pt",
    },
    "yolov8x-seg  (137.40)": {
        "model_url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8x-seg.pt",
    },
}

labelName = {
    0: "person",  
    1: "bicycle",  
    2: "car", 
    3: "motorcycle",  
    4: "airplane", 
    5: "bus",  
    6: "train", 
    7: "truck",  
    8: "boat", 
    9: "traffic light",  
    10: "fire hydrant",  
    11: "stop sign", 
    12: "parking meter", 
    13: "bench", 
    14: "bird",  
    15: "cat",  
    16: "dog",  
    17: "horse",  
    18: "sheep",  
    19: "cow",  
    20: "elephant", 
    21: "bear",  
    22: "zebra",  
    23: "giraffe",  
    24: "backpack",  
    25: "umbrella",  
    26: "handbag",  
    27: "tie",  
    28: "suitcase",  
    29: "frisbee",  
    30: "skis",  
    31: "snowboard",  
    32: "sports ball",  
    33: "kite",  
    34: "baseball bat",  
    35: "baseball glove",  
    36: "skateboard",  
    37: "surfboard", 
    38: "tennis racket",  
    39: "bottle",  
    40: "wine glass",  
    41: "cup",  
    42: "fork",  
    43: "knife",  
    44: "spoon",  
    45: "bowl",  
    46: "banana",  
    47: "apple",  
    48: "sandwich",  
    49: "orange",  
    50: "broccoli",  
    51: "carrot", 
    52: "hot dog", 
    53: "pizza",  
    54: "donut",  
    55: "cake",  
    56: "chair",  
    57: "couch",  
    58: "potted plant",  
    59: "bed",  
    60: "dining table",  
    61: "toilet",  
    62: "tv", 
    63: "laptop",  
    64: "mouse",  
    65: "remote",  
    66: "keyboard",  
    67: "cell phone",  
    68: "microwave",  
    69: "oven", 
    70: "toaster", 
    71: "sink", 
    72: "refrigerator",  
    73: "book",
    74: "clock", 
    75: "vase",  
    76: "scissors", 
    77: "teddy bear",  
    78: "hair drier",  
    79: "toothbrush",
}

def get_local_filepath(url, dirname, local_file_name=None):
    if not local_file_name:
        parsed_url = urlparse(url)
        local_file_name = os.path.basename(parsed_url.path)

    destination = folder_paths.get_full_path(dirname, local_file_name)
    if destination:
        logger.warn(f"using extra model: {destination}")
        return destination

    folder = os.path.join(folder_paths.models_dir, dirname)
    if not os.path.exists(folder):
        os.makedirs(folder)

    destination = os.path.join(folder, local_file_name)
    if not os.path.exists(destination):
        logger.warn(f"downloading {url} to {destination}")
        download_url_to_file(url, destination)
    return destination

def get_classes(label):
    label = label.lower()
    labels = label.split(",")
    result = []
    for l in labels:
        for key, value in labelName.items():
            if l == value:
                result.append(key)
                break
    return result

def get_yolov8_label_list():
    result = []
    for key, value in labelName.items():
        result.append(value)
    return result


def get_model_list():
    input_dir = folder_paths.get_input_directory()
    files = []
    for f in os.listdir(input_dir):
        if os.path.isfile(os.path.join(input_dir, f)):
            file_parts = f.split('.')
            if len(file_parts) > 1 and (file_parts[-1] == "pt"):
                files.append(f)
    return sorted(files)

def list_yolov8_model():
    return list(yolov8_model_list.keys())

def load_yolov8_model(model_name):
    yolov8_checkpoint_path = get_local_filepath(
        yolov8_model_list[model_name]["model_url"], yolov8_model_dir_name)
    model_file_name = os.path.basename(yolov8_checkpoint_path)
    model = YOLO(yolov8_checkpoint_path)
    return model

def load_yolov8_model_path(yolov8_checkpoint_path):
    model_file_name = os.path.basename(yolov8_checkpoint_path)
    model = YOLO(yolov8_checkpoint_path)
    return model

def is_url(url):
    return url.split("://")[0] in ["http", "https"]

def validate_path(path, allow_none=False, allow_url=True):
    if path is None:
        return allow_none
    if is_url(path):
        return True if allow_url else "URLs are unsupported for this path"
    if not os.path.isfile(path.strip("\"")):
        return "Invalid file path: {}".format(path)
    if not path.endswith('.pt'):
        return "Invalid file extension. Only .pt files are supported."
    return True

# modified from https://stackoverflow.com/questions/22058048/hashing-a-file-in-python
def calculate_file_hash(filename: str, hash_every_n: int = 1):
    # Larger video files were taking >.5 seconds to hash even when cached,
    # so instead the modified time from the filesystem is used as a hash
    h = hashlib.sha256()
    h.update(filename.encode())
    h.update(str(os.path.getmtime(filename)).encode())
    return h.hexdigest()

def yolov8_segment(model, image, label_name, threshold):
    image_tensor = image
    image_np = image_tensor.cpu().numpy()  # Change from CxHxW to HxWxC for Pillow
    image = Image.fromarray(
        (image_np.squeeze(0) * 255).astype(np.uint8)
    )  # Convert float [0,1] tensor to uint8 image

    if label_name is not None:
        classes = get_classes(label_name)
    else:
        classes = []
    results = model(image, classes=classes, conf=threshold)

    im_array = results[0].plot()  # plot a BGR numpy array of predictions
    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image

    image_tensor_out = torch.tensor(
        np.array(im).astype(np.float32) / 255.0
    )  # Convert back to CxHxW
    image_tensor_out = torch.unsqueeze(image_tensor_out, 0)

    res_mask=[]

    for result in results:
        masks = result.masks.data
        res_mask.append(torch.sum(masks, dim=0))
    return (image_tensor_out, res_mask)

def yolov8_detect(model, image, label_name, json_type, threshold):
    image_tensor = image
    image_np = image_tensor.cpu().numpy()  # Change from CxHxW to HxWxC for Pillow
    image = Image.fromarray(
        (image_np.squeeze(0) * 255).astype(np.uint8)
    )  # Convert float [0,1] tensor to uint8 image

    if label_name is not None:
        classes = get_classes(label_name)
    else:
        classes = []
    results = model(image, classes=classes, conf=threshold)

    im_array = results[0].plot()  # plot a BGR numpy array of predictions
    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image

    image_tensor_out = torch.tensor(
        np.array(im).astype(np.float32) / 255.0
    )  # Convert back to CxHxW
    image_tensor_out = torch.unsqueeze(image_tensor_out, 0)

    yolov8_json = []
    res_mask = []
    for result in results:
        labelme_data = {
            "version": "4.5.6",
            "flags": {},
            "shapes": [],
            "imagePath": result.path,
            "imageData": None,
            "imageHeight": result.orig_shape[0],
            "imageWidth": result.orig_shape[1],
        }
        for box in result.boxes:
            mask = np.zeros((result.orig_shape[0], result.orig_shape[1], 1), dtype=np.uint8)
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            label = labelName[int(box.cls)]
            points = [[x1, y1], [x2, y2]]
            shape = {
                "label": label,
                "points": points,
                "group_id": None,
                "shape_type": "rectangle",
                "flags": {},
            }
            json = [label, x1, y1, x2, y2]
            yolov8_json.append(json)
            labelme_data["shapes"].append(shape)
            cv2.rectangle(
                mask, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 255), -1
            )
            mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
            res_mask.append(mask_tensor)

    if json_type == "Labelme":
        json_data = labelme_data
    else:
        json_data = yolov8_json

    return (image_tensor_out, json_data, res_mask)


class LoadYolov8Model:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": (list_yolov8_model(),),
            },
        }
    CATEGORY = "Comfyui-Yolov8-JSON"
    FUNCTION = "main"
    RETURN_TYPES = ("YOLOV8_MODEL", )

    def main(self, model_name):
        model = load_yolov8_model(model_name)
        return (model,)

class LoadYolov8ModelFromPath:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_path": (
                    "STRING",
                    {"default": "/ComfyUI/models/yolov8/yolov8l.pt",}
                ),
            },
        }

    CATEGORY = "Comfyui-Yolov8-JSON"
    FUNCTION = "main"
    RETURN_TYPES = ("YOLOV8_MODEL",)

    def main(self, model_path):
        model_path = folder_paths.get_annotated_filepath(model_path.strip('"'))
        if model_path is None or validate_path(model_path) != True:
            raise Exception("model is not a valid path: " + model_path)
        model = load_yolov8_model_path(model_path)
        return (model,)

    @classmethod
    def IS_CHANGED(s, model_path):
        model_path = folder_paths.get_annotated_filepath(model_path)
        return calculate_file_hash(model_path)

class ApplyYolov8Model:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "yolov8_model": ("YOLOV8_MODEL", {}),
                "image": ("IMAGE",),
                "detect": (
                    ["all", "choose", "input"],
                    {"default": "all"},
                ),
                "label_name": (
                    "STRING",
                    {"default": "person,cat,dog", "multiline": False},
                ),
                "label_list": (
                    get_yolov8_label_list(),
                    {"default": "person"},
                ),
                "json_type": (
                    ["Labelme", "yolov8"],
                    {"default": "Labelme"},
                ),
                "threshold": (
                    "FLOAT",
                    {"default": 0.25, "min": 0.01, "max": 1.0, "step": 0.01},
                ),
            },
        }

    CATEGORY = "Comfyui-Yolov8-JSON"
    FUNCTION = "main"
    RETURN_TYPES = ("IMAGE", "JSON", "MASK")

    def main(self, yolov8_model, image, detect , label_name,label_list,json_type, threshold):
        res_images = []
        res_jsons = []
        res_masks = []
        for item in image:
            # Check and adjust image dimensions if needed
            if len(item.shape) == 3:
                item = item.unsqueeze(0)  # Add a batch dimension if missing

            label=None
            if(detect == "choose"):
                label=label_list
            else:
                label=label_name

            image_out, json, masks = yolov8_detect(
                yolov8_model, item, label, json_type, threshold
            )
            res_images.append(image_out)
            res_jsons.append(json)
            res_masks.extend(masks)
        return (torch.cat(res_images, dim=0), res_jsons, torch.cat(res_masks, dim=0))


class ApplyYolov8ModelSeg:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "yolov8_model": ("YOLOV8_MODEL", {}),
                "image": ("IMAGE",),
                "detect": (
                    ["all", "choose", "input"],
                    {"default": "all"},
                ),
                "label_name": (
                    "STRING",
                    {"default": "person,cat,dog", "multiline": False},
                ),
                "label_list": (
                    get_yolov8_label_list(),
                    {"default": "person"},
                ),
                "threshold": (
                    "FLOAT",
                    {"default": 0.25, "min": 0.01, "max": 1.0, "step": 0.01},
                ),
            },
        }

    CATEGORY = "Comfyui-Yolov8-JSON"
    FUNCTION = "main"
    RETURN_TYPES = ("IMAGE", "MASK")

    def main(
        self, yolov8_model, image, detect, label_name, label_list, threshold
    ):
        res_images = []
        res_masks = []
        for item in image:
            # Check and adjust image dimensions if needed
            if len(item.shape) == 3:
                item = item.unsqueeze(0)  # Add a batch dimension if missing

            label = None
            if detect == "choose":
                label = label_list
            else:
                label = label_name

            image_out,  masks = yolov8_segment(yolov8_model, item, label, threshold)
            res_images.append(image_out)
            res_masks.extend(masks)
        return (torch.cat(res_images, dim=0), torch.cat(res_masks, dim=0))


class SaveLabelmeJson:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {}),
                "labelme_json": ("JSON", {}),
                "folder_name": (
                    "STRING",
                    {"default": "GroundingDino", "multiline": False},
                ),
                "filename_prefix": (
                    "STRING",
                    {"default": "GroundingDino", "multiline": False},
                ),
            }
        }

    CATEGORY = "Comfyui-Yolov8-JSON"
    FUNCTION = "main"
    RETURN_TYPES = ("STRING",)

    def main(self, image, labelme_json, folder_name, filename_prefix):

        if len(labelme_json) != len(image):
            return '0'

        array_length = len(labelme_json)
        num_digits = len(str(array_length))
        count = 0

        # get outpu folder
        folder = folder_paths.output_directory
        output_dir = os.path.join(folder, folder_name)

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        for item, label in zip(image, labelme_json):
            image_pil = Image.fromarray(
                np.clip(255.0 * item.cpu().numpy(), 0, 255).astype(np.uint8)
            ).convert("RGB")

            count_str = f"{count:0{num_digits}d}"
            file_name = filename_prefix + "_" + count_str

            # save_image
            image_path = os.path.join(output_dir, file_name + ".jpg")
            image_pil.save(image_path)

            # save_json
            label["imagePath"] = file_name + ".jpg"
            json_path = os.path.join(output_dir, file_name + ".json")
            with open(json_path, "w") as json_file:
                json.dump(label, json_file, indent=4)

            count += 1

        return str(count)


def checkLabel(label, show_prompt):
    label = label.lower().split("(")[0]
    labels = show_prompt.split(",")
    for l in labels:
        new_label = l.lower()
        if label == new_label:
            return True
    return False


def parse_json_string(json_string):
    try:
        return json.loads(json_string)
    except json.JSONDecodeError as e:
        print(f"{json_string} decode error: {e}")
        return None


def plot_boxes_to_image(
    image_pil, labelme_json, show_prompt, event_prompt, prompt_name, show_threshold
):
    H = labelme_json["imageHeight"]
    W = labelme_json["imageWidth"]
    shapes = labelme_json["shapes"]
    prompt_list = parse_json_string(prompt_name)

    res_mask = []
    res_image = []

    font_scale = 1
    box_color = (255, 0, 0)
    text_color = (255, 255, 255)

    # Convert the image to a PIL image for drawing text
    draw = ImageDraw.Draw(image_pil)

    # Load a TTF font file for drawing Chinese text
    current_file_path = os.path.dirname(os.path.abspath(__file__))
    font_path = os.path.join(current_file_path, "docs", "PingFangRegular.ttf") 
    font_size = 20
    font = ImageFont.truetype(font_path, font_size)

    labelme_data = {
        "version": "4.5.6",
        "flags": {},
        "shapes": [],
        "imagePath": None,
        "imageData": None,
        "imageHeight": H,
        "imageWidth": W,
    }

    for shape in shapes:
        label = shape["label"]
        points = shape["points"]
        [x1, y1], [x2, y2] = points

        # if label is not in show, do not draw the label
        if show_prompt != "all" and show_prompt != "":
            if not checkLabel(label, show_prompt):
                continue

        # if label is event, color is red, else color is green
        if event_prompt != "all" and event_prompt != "":
            if checkLabel(label, event_prompt):
                box_color = (255, 0, 0)
                text_color = (255, 255, 255)
            else:
                box_color = (0, 255, 0)
                text_color = (255, 255, 255)

        # change label
        if prompt_list is not None and label in prompt_list:
            label = prompt_list[label]

        if "threshold" in shape and show_threshold=='yes':
            label = label + ":" + shape["threshold"]

        labelme_data["shapes"].append(shape)

        # Draw rectangle on the copied image
        draw.rectangle([(x1, y1), (x2, y2)], outline=box_color, width=3)

        # Get text size using textbbox
        text_bbox = draw.textbbox((0, 0), label, font=font)
        text_size = (text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1])

        label_ymin = max(y1, text_size[1] + 10)
        draw.rectangle([(x1, y1 - text_size[1] - 10), (x1 + text_size[0], y1)], fill=box_color)
        draw.text((x1, y1 - text_size[1] - 10), label, font=font, fill=text_color)

        # Draw mask
        mask = np.zeros((H, W, 1), dtype=np.uint8)
        cv2.rectangle(mask, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 255), -1)
        mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
        res_mask.append(mask_tensor)

    if len(res_mask) == 0:
        mask = np.zeros((H, W, 1), dtype=np.uint8)
        mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
        res_mask.append(mask_tensor)

    # Convert the modified image to a torch tensor
    image_with_boxes = np.array(image_pil)

    image_with_boxes_tensor = torch.from_numpy(image_with_boxes.astype(np.float32) / 255.0)
    image_with_boxes_tensor = torch.unsqueeze(image_with_boxes_tensor, 0)
    res_image.append(image_with_boxes_tensor)

    return res_image, res_mask, labelme_data


class DrawLabelmeJson:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {}),
                "labelme_json": ("JSON", {}),
                "show_prompt": (
                    "STRING",
                    {"default": "all", "multiline": False},
                ),
                "event_prompt": (
                    "STRING",
                    {"default": "all", "multiline": False},
                ),
                "prompt_name": (
                    "STRING",
                    {
                        "default": '{"head":"no helmet","helmet":"helmet"}',
                        "multiline": False,
                    },
                ),
                "show_threshold": (
                    ["yes", "no"],
                    {"default": "no"},
                ),
            }
        }

    CATEGORY = "Comfyui-Yolov8-JSON"
    FUNCTION = "main"
    RETURN_TYPES = (
        "IMAGE",
        "MASK",
        "JSON",
    )

    def main(
        self,
        image,
        labelme_json,
        show_prompt,
        event_prompt,
        prompt_name,
        show_threshold,
    ):

        res_images = []
        res_masks = []
        res_labels = []

        for item, labelme in zip(image, labelme_json):
            image_pil = Image.fromarray(np.clip(255.0 * item.cpu().numpy(), 0, 255).astype(np.uint8)).convert("RGB")
            image_tensor, mask_tensor, labelme_data = plot_boxes_to_image(
                image_pil,
                labelme,
                show_prompt,
                event_prompt,
                prompt_name,
                show_threshold,
            )
            res_images.extend(image_tensor)
            res_masks.extend(mask_tensor)
            res_labels.append(labelme_data)

            if len(res_images) == 0:
                res_images.extend(item)
            if len(res_masks) == 0:
                mask = np.zeros((height, width, 1), dtype=np.uint8)
                empty_mask = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
                res_masks.extend(empty_mask)

        return (torch.cat(res_images, dim=0), torch.cat(res_masks, dim=0), res_labels)