"""Commands to manage datasets local datasets and export between platforms
(Label Studio, HuggingFace Hub, local dataset,...)."""

import json
import random
import shutil
import typing
from pathlib import Path
from typing import Annotated, Optional

import typer
from openfoodfacts import Flavor
from openfoodfacts.utils import get_logger

from .label_studio import typer_description
from ..config import config, check_required_field
from ..types import ExportDestination, ExportSource, TaskType

app = typer.Typer(no_args_is_help=True)

logger = get_logger(__name__)


@app.command()
def check(
    dataset_dir: Annotated[
        Path,
        typer.Option(
            help="Path to the dataset directory", exists=True, file_okay=False
        ),
    ],
    remove: Annotated[
        bool,
        typer.Option(help="Remove duplicate images from the dataset"),
    ] = False,
):
    """Check a local dataset in Ultralytics format for duplicate images."""

    from ..check import check_local_dataset

    check_local_dataset(dataset_dir, remove=remove)


@app.command()
def split_train_test(
    task_type: TaskType, dataset_dir: Path, output_dir: Path, train_ratio: float = 0.8
):
    """Split a local dataset into training and test sets.

    Only classification tasks are supported.
    """
    if task_type == TaskType.classification:
        class_dirs = [d for d in dataset_dir.iterdir() if d.is_dir()]
        logger.info("Found classes: %s", [d.name for d in class_dirs])

        output_dir.mkdir(parents=True, exist_ok=True)
        train_dir = output_dir / "train"
        test_dir = output_dir / "test"
        train_dir.mkdir(parents=True, exist_ok=True)
        test_dir.mkdir(parents=True, exist_ok=True)

        for class_dir in class_dirs:
            input_paths = list(class_dir.glob("*"))
            random.shuffle(input_paths)

            test_count = int(len(input_paths) * (1 - train_ratio))
            if test_count == 0:
                logger.warning("Not enough samples, skipping class: %s", class_dir.name)
                continue

            test_paths = input_paths[:test_count]
            train_paths = input_paths[test_count:]

            for output_dir, input_paths in (
                (train_dir, train_paths),
                (test_dir, test_paths),
            ):
                output_cls_dir = output_dir / class_dir.name
                output_cls_dir.mkdir(parents=True, exist_ok=True)

                for path in input_paths:
                    logger.info("Copying: %s to %s", path, output_cls_dir)
                    shutil.copy(path, output_cls_dir / path.name)
    else:
        raise typer.BadParameter("Unsupported task type")


@app.command()
def convert_object_detection_dataset(
    repo_id: Annotated[
        str, typer.Option(help="Hugging Face Datasets repository ID to convert")
    ],
    output_file: Annotated[
        Path, typer.Option(help="Path to the output JSON file", exists=False)
    ],
):
    """Convert object detection dataset from Hugging Face Datasets to Label
    Studio format, and save it to a JSON file."""
    from datasets import load_dataset

    from labelr.sample.object_detection import (
        format_object_detection_sample_from_hf_to_ls,
    )

    logger.info("Loading dataset: %s", repo_id)
    ds = load_dataset(repo_id)
    logger.info("Dataset loaded: %s", tuple(ds.keys()))

    with output_file.open("wt") as f:
        for split in ds.keys():
            logger.info("Processing split: %s", split)
            for sample in ds[split]:
                label_studio_sample = format_object_detection_sample_from_hf_to_ls(
                    sample, split=split
                )
                f.write(json.dumps(label_studio_sample) + "\n")


@app.command()
def export(
    from_: Annotated[ExportSource, typer.Option("--from", help="Input source to use.")],
    to: Annotated[ExportDestination, typer.Option(help="Where to export the data.")],
    api_key: Annotated[
        str | None, typer.Option(help=typer_description.LABEL_STUDIO_API_KEY)
    ] = config.label_studio_api_key,
    task_type: Annotated[
        TaskType, typer.Option(help="Type of task to export.")
    ] = TaskType.object_detection,
    repo_id: Annotated[
        str | None,
        typer.Option(
            help="Hugging Face Datasets repository ID to convert (only if --from or --to is `hf`)."
        ),
    ] = None,
    label_names: Annotated[
        str | None,
        typer.Option(help="Label names to use, as a comma-separated list"),
    ] = None,
    skip_labels: Annotated[
        str | None,
        typer.Option(
            help="Label names to skip  as a comma-separated list. If the label of an "
            "annotation result is in this list, it will be skipped. Only used if the "
            "source is Label Studio and the task type is object detection."
        ),
    ] = None,
    project_id: Annotated[
        int | None, typer.Option(help=typer_description.LABEL_STUDIO_PROJECT_ID)
    ] = config.label_studio_project_id,
    view_id: Annotated[
        int | None,
        typer.Option(
            help="ID of the Label Studio view, if any. This option is useful "
            "to filter the task to export."
        ),
    ] = None,
    label_studio_url: Annotated[
        str, typer.Option(help=typer_description.LABEL_STUDIO_URL)
    ] = config.label_studio_url,
    output_dir: Annotated[
        Path | None,
        typer.Option(
            help="Path to the output directory. Only used if the destintation (`to`) is `ultralytics`",
            file_okay=False,
        ),
    ] = None,
    dataset_dir: Annotated[
        Path | None,
        typer.Option(help="Path to the dataset directory, only for Ultralytics source"),
    ] = None,
    download_images: Annotated[
        bool,
        typer.Option(
            help="if True, don't use HF images and download images from the server"
        ),
    ] = False,
    is_openfoodfacts_dataset: Annotated[
        bool,
        typer.Option(
            help="Whether the Ultralytics dataset is an Open Food Facts dataset, only "
            "for Ultralytics source. This is used:\n"
            "- to generate the correct image URLs from each image name, when exporting "
            "from Ultralytics to Hugging Face Datasets.\n"
            "- to include additional metadata fields specific to Open Food Facts "
            "(`barcode` and `off_image_id`) when exporting from Label Studio to "
            "Hugging Face Datasets."
        ),
    ] = False,
    openfoodfacts_flavor: Annotated[
        Flavor,
        typer.Option(
            help="Flavor of the Open Food Facts dataset to use for image URLs, only "
            "for Ultralytics source if is_openfoodfacts_dataset is True. This is used to "
            "generate the correct image URLs each image name. This option is ignored if "
            "is_openfoodfacts_dataset is False."
        ),
    ] = Flavor.off,
    train_ratio: Annotated[
        float,
        typer.Option(
            help="Train ratio for splitting the dataset, if the split name is not "
            "provided. Only used if the source is Label Studio and the destination "
            "is Ultralytics."
        ),
    ] = 0.8,
    image_max_size: Annotated[
        int | None,
        typer.Option(
            help="Maximum size (in pixels) for the images. If None, no resizing is performed."
            "Otherwise, the longest side of the image will be resized to this value, "
            "keeping the aspect ratio."
        ),
    ] = None,
    error_raise: Annotated[
        bool,
        typer.Option(
            help="Raise an error if an image download fails, only for Ultralytics"
        ),
    ] = True,
    use_aws_cache: Annotated[
        bool,
        typer.Option(
            help="Use the AWS S3 cache for image downloads instead of images.openfoodfacts.org, "
            "it is ignored if the export format is not Ultralytics"
        ),
    ] = False,
    merge_labels: Annotated[
        bool,
        typer.Option(help="Merge multiple labels into a single label"),
    ] = False,
    revision: Annotated[
        str,
        typer.Option(
            help="Revision (branch, tag or commit) for the Hugging Face Datasets repository. "
            "Only used when importing from or exporting to Hugging Face Datasets."
        ),
    ] = "main",
):
    """Export Label Studio annotation, either to Hugging Face Datasets or
    local files (ultralytics format)."""
    from label_studio_sdk.client import LabelStudio

    from labelr.export.object_detection import (
        export_from_hf_to_ultralytics_object_detection,
    )
    from labelr.export.object_detection import (
        export_from_ls_to_hf_object_detection,
        export_from_ls_to_ultralytics_object_detection,
    )
    from labelr.export.common import export_from_ultralytics_to_hf

    check_required_field("--project-id", project_id)

    if (to == ExportDestination.hf or from_ == ExportSource.hf) and repo_id is None:
        raise typer.BadParameter("Repository ID is required for export/import with HF")

    if from_ == ExportSource.ultralytics and dataset_dir is None:
        raise typer.BadParameter(
            "Dataset directory is required for export from Ultralytics source"
        )

    label_names_list: list[str] | None = None

    if label_names is None:
        if to == ExportDestination.hf:
            raise typer.BadParameter("Label names are required for HF export")
        if from_ == ExportSource.ls:
            raise typer.BadParameter(
                "Label names are required for export from LS source"
            )
    else:
        label_names = typing.cast(str, label_names)
        label_names_list = label_names.split(",")

    skip_labels_list = []
    if skip_labels is not None:
        skip_labels = typing.cast(str, skip_labels)
        skip_labels_list = skip_labels.split(",")

    if from_ == ExportSource.ls:
        if project_id is None:
            raise typer.BadParameter("Project ID is required for LS export")
        if api_key is None:
            raise typer.BadParameter("API key is required for LS export")

    if to == ExportDestination.ultralytics and output_dir is None:
        raise typer.BadParameter("Output directory is required for Ultralytics export")

    if from_ == ExportSource.ls:
        if task_type != TaskType.object_detection:
            raise typer.BadParameter(
                "Only object detection task is currently supported with LS source"
            )
        ls = LabelStudio(base_url=label_studio_url, api_key=api_key)
        if to == ExportDestination.hf:
            repo_id = typing.cast(str, repo_id)
            export_from_ls_to_hf_object_detection(
                ls,
                repo_id=repo_id,
                label_names=typing.cast(list[str], label_names_list),
                project_id=typing.cast(int, project_id),
                is_openfoodfacts_dataset=is_openfoodfacts_dataset,
                merge_labels=merge_labels,
                use_aws_cache=use_aws_cache,
                revision=revision,
                view_id=view_id,
                image_max_size=image_max_size,
                skip_labels=skip_labels_list,
            )
        elif to == ExportDestination.ultralytics:
            export_from_ls_to_ultralytics_object_detection(
                ls,
                typing.cast(Path, output_dir),
                typing.cast(list[str], label_names_list),
                typing.cast(int, project_id),
                train_ratio=train_ratio,
                error_raise=error_raise,
                merge_labels=merge_labels,
                use_aws_cache=use_aws_cache,
                view_id=view_id,
                image_max_size=image_max_size,
                skip_labels=skip_labels_list,
            )

    elif from_ == ExportSource.hf:
        if task_type != TaskType.object_detection:
            raise typer.BadParameter(
                "Only object detection task is currently supported with HF source"
            )
        if to == ExportDestination.ultralytics:
            export_from_hf_to_ultralytics_object_detection(
                typing.cast(str, repo_id),
                typing.cast(Path, output_dir),
                download_images=download_images,
                error_raise=error_raise,
                use_aws_cache=use_aws_cache,
                revision=revision,
                image_max_size=image_max_size,
                skip_labels=skip_labels_list,
            )
        else:
            raise typer.BadParameter("Unsupported export format")
    elif from_ == ExportSource.ultralytics:
        if task_type != TaskType.classification:
            raise typer.BadParameter(
                "Only classification task is currently supported with Ultralytics source"
            )
        if to == ExportDestination.hf:
            export_from_ultralytics_to_hf(
                task_type=task_type,
                dataset_dir=typing.cast(Path, dataset_dir),
                repo_id=typing.cast(str, repo_id),
                merge_labels=merge_labels,
                label_names=typing.cast(list[str], label_names_list),
                is_openfoodfacts_dataset=is_openfoodfacts_dataset,
                openfoodfacts_flavor=openfoodfacts_flavor,
            )


@app.command()
def export_llm_ds(
    dataset_path: Annotated[
        Path, typer.Option(..., help="Path to the JSONL dataset file")
    ],
    repo_id: Annotated[
        str, typer.Option(..., help="Hugging Face Datasets repository ID to export to")
    ],
    split: Annotated[str, typer.Option(..., help="Dataset split to export")],
    revision: Annotated[
        str,
        typer.Option(
            help="Revision (branch, tag or commit) for the Hugging Face Datasets repository."
        ),
    ] = "main",
    tmp_dir: Annotated[
        Path | None,
        typer.Option(
            help="Path to the temporary directory used to store intermediate sample files "
            "created when building the HF dataset.",
        ),
    ] = None,
    image_max_size: Annotated[
        int | None,
        typer.Option(
            help="Maximum size (in pixels) for the images. If None, no resizing is performed.",
        ),
    ] = None,
):
    """Export LLM image extraction dataset with images only to Hugging Face
    Datasets.
    """
    from labelr.export.llm import export_to_hf_llm_image_extraction
    from labelr.sample.llm import load_llm_image_extraction_dataset_from_jsonl

    sample_iter = load_llm_image_extraction_dataset_from_jsonl(
        dataset_path=dataset_path
    )
    export_to_hf_llm_image_extraction(
        sample_iter,
        split=split,
        repo_id=repo_id,
        revision=revision,
        tmp_dir=tmp_dir,
        image_max_size=image_max_size,
    )


@app.command()
def update_llm_ds(
    dataset_path: Annotated[
        Path, typer.Option(help="Path to the JSONL containing the updates.")
    ],
    repo_id: Annotated[
        str, typer.Option(help="Hugging Face Datasets repository ID to update")
    ],
    split: Annotated[str, typer.Option(help="Dataset split to use")],
    revision: Annotated[
        str,
        typer.Option(
            help="Revision (branch, tag or commit) to use when pushing the new version "
            "of the Hugging Face Dataset."
        ),
    ] = "main",
    tmp_dir: Annotated[
        Path | None,
        typer.Option(
            help="Path to a temporary directory to use for image processing",
        ),
    ] = None,
    show_diff: Annotated[
        bool,
        typer.Option(
            help="Show the differences between the original sample and the update. If "
            "True, the updated dataset is not pushed to the Hub. Useful to review the "
            "updates before applying them.",
        ),
    ] = False,
):
    """Update an existing LLM image extraction dataset, by updating the
    `output` field of each sample in the dataset.

    The `--dataset_path` JSONL file should contain items with two fields:

    - `image_id`: The image ID of the sample to update in the Hugging Face
        dataset.
    - `output`: The new output data to set for the sample.
    """
    import sys
    from difflib import Differ

    import orjson
    from datasets import load_dataset
    from diskcache import Cache

    dataset = load_dataset(repo_id, split=split)

    # Populate cache with the updates
    cache = Cache(directory=tmp_dir or None)
    with dataset_path.open("r") as f:
        for line in map(orjson.loads, f):
            if "image_id" not in line or "output" not in line:
                raise ValueError(
                    "Each item in the update JSONL file must contain `image_id` and `output` fields"
                )
            image_id = line["image_id"]
            output = line["output"]

            if not isinstance(output, str):
                output = orjson.dumps(output).decode("utf-8")

            cache[image_id] = output

    def apply_updates(sample):
        image_id = sample["image_id"]
        if image_id in cache:
            cached_item = cache[image_id]
            sample["output"] = cached_item
        return sample

    if show_diff:
        differ = Differ()
        for sample in dataset:
            image_id = sample["image_id"]
            if image_id in cache:
                cached_item = orjson.loads(cache[image_id])
                original_item = orjson.loads(sample["output"])
                cached_item_str = orjson.dumps(
                    cached_item, option=orjson.OPT_INDENT_2
                ).decode("utf8")
                original_item_str = orjson.dumps(
                    original_item, option=orjson.OPT_INDENT_2
                ).decode("utf8")
                diff = list(
                    differ.compare(
                        original_item_str.splitlines(keepends=True),
                        cached_item_str.splitlines(keepends=True),
                    )
                )
                sys.stdout.writelines(diff)
                sys.stdout.write("\n" + "-" * 30 + "\n")

    else:
        updated_dataset = dataset.map(apply_updates, batched=False)
        updated_dataset.push_to_hub(repo_id, split=split, revision=revision)
