import functools
import logging
import pickle
import random
import tempfile
import typing
from pathlib import Path

import datasets
import tqdm
from label_studio_sdk.client import LabelStudio
from openfoodfacts.images import download_image
from PIL import Image, ImageOps

from labelr.export.common import _pickle_sample_generator
from labelr.sample.object_detection import (
    format_object_detection_sample_to_hf,
    get_hf_object_detection_features,
)

logger = logging.getLogger(__name__)


def export_from_ls_to_hf_object_detection(
    ls: LabelStudio,
    repo_id: str,
    label_names: list[str],
    project_id: int,
    is_openfoodfacts_dataset: bool,
    image_max_size: int | None = None,
    view_id: int | None = None,
    merge_labels: bool = False,
    use_aws_cache: bool = True,
    revision: str = "main",
    skip_labels: list[str] | None = None,
) -> None:
    """Export annotations from a Label Studio project to a Hugging Face
    dataset.

    The Label Studio project should be an object detection project.

    Args:
        ls (LabelStudio): Label Studio client instance.
        repo_id (str): Hugging Face repository ID to push the dataset to.
        label_names (list[str]): List of label names in the project.
        project_id (int): Label Studio project ID to export from.
        is_openfoodfacts_dataset (bool): Whether the dataset is an Open Food
            Facts dataset. If True, the dataset will include additional
            metadata fields specific to Open Food Facts (`barcode` and
            `off_image_id`).
        image_max_size (int | None): Maximum size (in pixels) for the images.
            If None, no resizing is performed. Defaults to None.
        view_id (int | None): Label Studio view ID to export from. If None,
            all tasks are exported. Defaults to None.
        merge_labels (bool): Whether to merge all labels into a single label
            named "object". Defaults to False.
        use_aws_cache (bool): Whether to use the AWS image cache when
            downloading images. Defaults to True.
        revision (str): The dataset revision to push to. Defaults to 'main'.
        skip_labels (list[str] | None): List of label names to skip. If the
            label of an annotation result is in this list, it will be skipped.
            Defaults to None.
    """
    if merge_labels:
        label_names = ["object"]

    logger.info(
        "Project ID: %d, label names: %s, skip labels: %s, repo_id: %s, revision: %s, view ID: %s",
        project_id,
        label_names,
        skip_labels,
        repo_id,
        revision,
        view_id,
    )

    for split in ["train", "val"]:
        logger.info("Processing split: %s", split)

        with tempfile.TemporaryDirectory() as tmp_dir_str:
            tmp_dir = Path(tmp_dir_str)
            logger.info("Saving samples to temporary directory: %s", tmp_dir)
            for i, task in tqdm.tqdm(
                enumerate(
                    ls.tasks.list(project=project_id, fields="all", view=view_id)
                ),
                desc="tasks",
            ):
                if task.data["split"] != split:
                    continue
                sample = format_object_detection_sample_to_hf(
                    task_data=task.data,
                    annotations=task.annotations,
                    label_names=label_names,
                    merge_labels=merge_labels,
                    use_aws_cache=use_aws_cache,
                    image_max_size=image_max_size,
                    skip_labels=skip_labels,
                )
                if sample is not None:
                    # Save output as pickle
                    with open(tmp_dir / f"{split}_{i:05}.pkl", "wb") as f:
                        pickle.dump(sample, f)

            features = get_hf_object_detection_features(is_openfoodfacts_dataset)
            hf_ds = datasets.Dataset.from_generator(
                functools.partial(_pickle_sample_generator, tmp_dir),
                features=features,
            )
            hf_ds.push_to_hub(repo_id, split=split, revision=revision)


def export_from_ls_to_ultralytics_object_detection(
    ls: LabelStudio,
    output_dir: Path,
    label_names: list[str],
    project_id: int,
    train_ratio: float = 0.8,
    error_raise: bool = True,
    merge_labels: bool = False,
    use_aws_cache: bool = True,
    view_id: int | None = None,
    image_max_size: int | None = None,
    skip_labels: list[str] | None = None,
):
    """Export annotations from a Label Studio project to the Ultralytics
    format.

    The Label Studio project should be an object detection project with a
    single rectanglelabels annotation result per task.

    Args:
        ls (LabelStudio): Label Studio client instance.
        output_dir (Path): Path to the output directory.
        label_names (list[str]): List of label names in the project.
        project_id (int): Label Studio project ID to export from.
        train_ratio (float): Ratio of training samples. The rest will be used
            for validation. Defaults to 0.8.
        error_raise (bool): Whether to raise an error if an image fails to
            download. If False, the image will be skipped. Defaults to True.
        merge_labels (bool): Whether to merge all labels into a single label
            named "object". Defaults to False.
        use_aws_cache (bool): Whether to use the AWS image cache when
            downloading images. Defaults to True.
        view_id (int | None): Label Studio view ID to export from. If None,
            all tasks are exported. Defaults to None.
        image_max_size (int | None): Maximum size (in pixels) for the images.
            If None, no resizing is performed. Defaults to None.
        skip_labels (list[str] | None): List of label names to skip. If the
            label of an annotation result is in this list, it will be skipped.
            Defaults to None.
    """
    if merge_labels:
        label_names = ["object"]

    skip_labels = skip_labels or []
    logger.info(
        "Project ID: %d, label names: %s, skip labels: %s",
        project_id,
        label_names,
        skip_labels,
    )

    data_dir = output_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    split_warning_displayed = False

    # NOTE: before, all images were sent to val, the last split
    label_dir = data_dir / "labels"
    images_dir = data_dir / "images"
    for split in ["train", "val"]:
        (label_dir / split).mkdir(parents=True, exist_ok=True)
        (images_dir / split).mkdir(parents=True, exist_ok=True)

    for task in tqdm.tqdm(
        ls.tasks.list(project=project_id, fields="all", view=view_id),
        desc="tasks",
    ):
        split = task.data.get("split")

        if split is None:
            if not split_warning_displayed:
                logger.warning(
                    "Split information not found, assigning randomly. "
                    "To avoid this, set the `split` field in the task data."
                )
                split_warning_displayed = True
            split = "train" if random.random() < train_ratio else "val"

        elif split not in ["train", "val"]:
            raise ValueError("Invalid split name: %s", split)

        if len(task.annotations) > 1:
            logger.warning("More than one annotation found, skipping")
            continue
        elif len(task.annotations) == 0:
            logger.debug("No annotation found, skipping")
            continue

        annotation = task.annotations[0]
        if annotation["was_cancelled"] is True:
            logger.debug("Annotation was cancelled, skipping")
            continue

        if "image_id" not in task.data:
            raise ValueError(
                "`image_id` field not found in task data. "
                "Make sure the task data contains the `image_id` "
                "field, which should be a unique identifier for the image."
            )
        if "image_url" not in task.data:
            raise ValueError(
                "`image_url` field not found in task data. "
                "Make sure the task data contains the `image_url` "
                "field, which should be the URL of the image."
            )
        image_id = task.data["image_id"]
        image_url = task.data["image_url"]

        has_valid_annotation = False
        with (label_dir / split / f"{image_id}.txt").open("w") as f:
            if not any(
                annotation_result["type"] == "rectanglelabels"
                for annotation_result in annotation["result"]
            ):
                continue

            for annotation_result in annotation["result"]:
                if annotation_result["type"] == "rectanglelabels":
                    value = annotation_result["value"]
                    x_min = value["x"] / 100
                    y_min = value["y"] / 100
                    width = value["width"] / 100
                    height = value["height"] / 100
                    label_name = (
                        label_names[0] if merge_labels else value["rectanglelabels"][0]
                    )
                    if label_name in skip_labels:
                        continue
                    label_id = (
                        label_names.index(label_name)
                        if label_name in label_names
                        else -1
                    )
                    if label_id == -1:
                        logger.warning(
                            "Label name '%s' not found in label_names, skipping annotation result: '%s'",
                            label_name,
                            annotation_result,
                        )
                        continue

                    # Save the labels in the Ultralytics format:
                    # - one label per line
                    # - each line is a list of 5 elements:
                    #   - label_id
                    #   - x_center
                    #   - y_center
                    #   - width
                    #   - height
                    x_center = x_min + width / 2
                    y_center = y_min + height / 2
                    f.write(f"{label_id} {x_center} {y_center} {width} {height}\n")
                    has_valid_annotation = True

        if has_valid_annotation:
            image = download_image(
                image_url,
                return_struct=False,
                error_raise=error_raise,
                use_cache=use_aws_cache,
            )
            if image is None:
                logger.error("Failed to download image: %s", image_url)
                continue

            image = typing.cast(Image.Image, image)

            # Rotate image according to exif orientation using Pillow
            ImageOps.exif_transpose(image, in_place=True)
            # Resize image if larger than max size
            if image_max_size is not None and (
                image.width > image_max_size or image.height > image_max_size
            ):
                image.thumbnail(
                    (image_max_size, image_max_size), Image.Resampling.LANCZOS
                )
            image.save(images_dir / split / f"{image_id}.jpg", format="JPEG")

    with (output_dir / "data.yaml").open("w") as f:
        f.write("path: data\n")
        f.write("train: images/train\n")
        f.write("val: images/val\n")
        f.write("test:\n")
        f.write("names:\n")
        for i, label_name in enumerate(label_names):
            f.write(f"  {i}: {label_name}\n")


def export_from_hf_to_ultralytics_object_detection(
    repo_id: str,
    output_dir: Path,
    download_images: bool = True,
    error_raise: bool = True,
    use_aws_cache: bool = True,
    image_max_size: int | None = None,
    revision: str = "main",
    skip_labels: list[str] | None = None,
):
    """Export annotations from a Hugging Face dataset project to the
    Ultralytics format.

    The Label Studio project should be an object detection project with a
    single rectanglelabels annotation result per task.

    Args:
        repo_id (str): Hugging Face repository ID to load the dataset from.
        output_dir (Path): Path to the output directory.
        download_images (bool): Whether to download images from URLs in the
            dataset. If False, the dataset is expected to contain an `image`
            field with the image data.
        error_raise (bool): Whether to raise an error if an image fails to
            download. If False, the image will be skipped. This option is only
            used if `download_images` is True. Defaults to True.
        use_aws_cache (bool): Whether to use the AWS image cache when
            downloading images. This option is only used if `download_images`
            is True. Defaults to True.
        image_max_size (int | None): Maximum size (in pixels) for the images.
            If None, no resizing is performed. Defaults to None.
        revision (str): The dataset revision to load. Defaults to 'main'.
        skip_labels (list[str] | None): List of label names to skip. If the
            label of an annotation result is in this list, it will be skipped.
            Defaults to None.
    """
    skip_labels = skip_labels or []
    logger.info(
        "Repo ID: %s, revision: %s, skip labels: %s", repo_id, revision, skip_labels
    )
    ds = datasets.load_dataset(repo_id, revision=revision)
    data_dir = output_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    label_id_to_name = {}

    split_map = {
        "train": "train",
        "val": "val",
    }
    if "val" not in ds and "test" in ds:
        logger.info("val split not found, using test split instead as val")
        split_map["val"] = "test"

    for split in ["train", "val"]:
        split_target = split_map[split]
        split_labels_dir = data_dir / "labels" / split
        split_labels_dir.mkdir(parents=True, exist_ok=True)
        split_images_dir = data_dir / "images" / split
        split_images_dir.mkdir(parents=True, exist_ok=True)

        for sample in tqdm.tqdm(ds[split_target], desc="samples"):
            image_id = sample["image_id"]

            if download_images:
                if "meta" not in sample or "image_url" not in sample["meta"]:
                    raise ValueError(
                        "`meta.image_url` field not found in sample. "
                        "Make sure the dataset contains the `meta.image_url` "
                        "field, which should be the URL of the image, or set "
                        "`download_images` to False."
                    )
                image_url = sample["meta"]["image_url"]
                image = download_image(
                    image_url,
                    return_struct=False,
                    error_raise=error_raise,
                    use_cache=use_aws_cache,
                )
                if image is None:
                    logger.error("Failed to download image: %s", image_url)
                    continue
            else:
                image = sample["image"]

            image = typing.cast(Image.Image, image)
            # Rotate image according to exif orientation using Pillow
            # If the image source is Hugging Face, EXIF data is not preserved,
            # so this step is only useful when downloading images.
            ImageOps.exif_transpose(image, in_place=True)
            if image.mode != "RGB":
                # Convert image to RGB if needed
                image = image.convert("RGB")
            # Resize image if larger than max size
            if image_max_size is not None and (
                image.width > image_max_size or image.height > image_max_size
            ):
                image.thumbnail(
                    (image_max_size, image_max_size), Image.Resampling.LANCZOS
                )
            image.save(split_images_dir / f"{image_id}.jpg")

            objects = sample["objects"]
            bboxes = objects["bbox"]
            label_ids = objects["category_id"]
            label_names = objects["category_name"]

            with (split_labels_dir / f"{image_id}.txt").open("w") as f:
                for bbox, label_id, label_name in zip(bboxes, label_ids, label_names):
                    if label_id not in label_id_to_name:
                        label_id_to_name[label_id] = label_name
                    if label_name in skip_labels:
                        continue
                    y_min, x_min, y_max, x_max = bbox
                    y_min = min(max(y_min, 0.0), 1.0)
                    x_min = min(max(x_min, 0.0), 1.0)
                    y_max = min(max(y_max, 0.0), 1.0)
                    x_max = min(max(x_max, 0.0), 1.0)
                    width = x_max - x_min
                    height = y_max - y_min
                    # Save the labels in the Ultralytics format:
                    # - one label per line
                    # - each line is a list of 5 elements:
                    #   - label_id
                    #   - x_center
                    #   - y_center
                    #   - width
                    #   - height
                    x_center = x_min + width / 2
                    y_center = y_min + height / 2
                    f.write(f"{label_id} {x_center} {y_center} {width} {height}\n")

    label_names = [x[1] for x in sorted(label_id_to_name.items(), key=lambda x: x[0])]
    with (output_dir / "data.yaml").open("w") as f:
        f.write("path: data\n")
        f.write("train: images/train\n")
        f.write("val: images/val\n")
        f.write("test:\n")
        f.write("names:\n")
        for i, label_name in enumerate(label_names):
            f.write(f"  {i}: {label_name}\n")
