from __future__ import annotations

from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Union

from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
from rich import print

repo_id = "Bingsu/adetailer"


@dataclass
class PredictOutput:
    bboxes: list[list[int | float]] = field(default_factory=list)
    masks: list[Image.Image] = field(default_factory=list)
    preview: Optional[Image.Image] = None


def hf_download(file: str):
    try:
        path = hf_hub_download(repo_id, file)
    except Exception:
        msg = f"[-] ADetailer: Failed to load model {file!r} from huggingface"
        print(msg)
        path = "INVALID"
    return path


def get_models(
    model_dir: Union[str, Path], huggingface: bool = True
) -> OrderedDict[str, Optional[str]]:
    model_dir = Path(model_dir)
    if model_dir.is_dir():
        model_paths = [
            p
            for p in model_dir.rglob("*")
            if p.is_file() and p.suffix in (".pt", ".pth")
        ]
    else:
        model_paths = []

    models = OrderedDict()
    if huggingface:
        models.update(
            {
                "face_yolov8n.pt": hf_download("face_yolov8n.pt"),
                "face_yolov8s.pt": hf_download("face_yolov8s.pt"),
                "hand_yolov8n.pt": hf_download("hand_yolov8n.pt"),
                "person_yolov8n-seg.pt": hf_download("person_yolov8n-seg.pt"),
                "person_yolov8s-seg.pt": hf_download("person_yolov8s-seg.pt"),
            }
        )
    models.update(
        {
            "mediapipe_face_full": None,
            "mediapipe_face_short": None,
            "mediapipe_face_mesh": None,
            "mediapipe_face_mesh_eyes_only": None,
        }
    )

    invalid_keys = [k for k, v in models.items() if v == "INVALID"]
    for key in invalid_keys:
        models.pop(key)

    for path in model_paths:
        if path.name in models:
            continue
        models[path.name] = str(path)

    return models


def create_mask_from_bbox(
    bboxes: list[list[float]], shape: tuple[int, int]
) -> list[Image.Image]:
    """
    Parameters
    ----------
        bboxes: list[list[float]]
            list of [x1, y1, x2, y2]
            bounding boxes
        shape: tuple[int, int]
            shape of the image (width, height)

    Returns
    -------
        masks: list[Image.Image]
        A list of masks

    """
    masks = []
    for bbox in bboxes:
        mask = Image.new("L", shape, 0)
        mask_draw = ImageDraw.Draw(mask)
        mask_draw.rectangle(bbox, fill=255)
        masks.append(mask)
    return masks


def create_bbox_from_mask(
    masks: list[Image.Image], shape: tuple[int, int]
) -> list[list[int]]:
    """
    Parameters
    ----------
        masks: list[Image.Image]
            A list of masks
        shape: tuple[int, int]
            shape of the image (width, height)

    Returns
    -------
        bboxes: list[list[float]]
        A list of bounding boxes

    """
    bboxes = []
    for mask in masks:
        mask = mask.resize(shape)
        bbox = mask.getbbox()
        if bbox is not None:
            bboxes.append(list(bbox))
    return bboxes