import json
import os

import datasets
import pandas as pd
import pytest
from huggingface_hub.errors import HfHubHTTPError

import mlflow.data
import mlflow.data.huggingface_dataset
from mlflow.data.code_dataset_source import CodeDatasetSource
from mlflow.data.dataset_source_registry import get_dataset_source_from_json
from mlflow.data.evaluation_dataset import EvaluationDataset
from mlflow.data.huggingface_dataset import HuggingFaceDataset
from mlflow.data.huggingface_dataset_source import HuggingFaceDatasetSource
from mlflow.exceptions import MlflowException
from mlflow.types.schema import Schema
from mlflow.types.utils import _infer_schema

from tests.helper_functions import skip_if_hf_hub_unhealthy
from tests.resources.data.dataset_source import SampleDatasetSource

pytestmark = skip_if_hf_hub_unhealthy()


@pytest.fixture(scope="module", autouse=True)
def prefetch_huggingface_datasets():
    """Pre-warm the HF cache so individual tests don't hit the Hub and risk HTTP 429s."""

    try:
        datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
        datasets.load_dataset(
            "cornell-movie-review-data/rotten_tomatoes",
            split="train",
            revision="aa13bc287fa6fcab6daf52f0dfb9994269ffea28",
            trust_remote_code=True,
        )
        datasets.load_dataset(
            "cornell-movie-review-data/rotten_tomatoes",
            split="train",
            revision="c33cbf965006dba64f134f7bef69c53d5d0d285d",
        )
        datasets.load_dataset(
            "fka/awesome-chatgpt-prompts",
            data_files={"train": "prompts.csv"},
            split="train",
        )
    except HfHubHTTPError as e:
        if e.response is not None and e.response.status_code == 429:
            pytest.skip(f"HF Hub returned 429 while pre-warming the cache: {e}")
        raise


def test_from_huggingface_dataset_constructs_expected_dataset():
    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
    mlflow_ds = mlflow.data.from_huggingface(ds, path="cornell-movie-review-data/rotten_tomatoes")

    assert isinstance(mlflow_ds, HuggingFaceDataset)
    assert mlflow_ds.ds == ds
    assert mlflow_ds.schema == _infer_schema(ds.to_pandas())
    assert mlflow_ds.profile == {
        "num_rows": ds.num_rows,
        "dataset_size": ds.dataset_size,
        "size_in_bytes": ds.size_in_bytes,
    }

    assert isinstance(mlflow_ds.source, HuggingFaceDatasetSource)

    with pytest.raises(KeyError, match="Found duplicated arguments*"):
        # Test that we raise an error if the same key is specified in both
        # `HuggingFaceDatasetSource` and `kwargs`.
        mlflow_ds.source.load(path="dummy_path")

    reloaded_ds = mlflow_ds.source.load()
    assert reloaded_ds.builder_name == ds.builder_name
    assert reloaded_ds.config_name == ds.config_name
    assert reloaded_ds.split == ds.split == "train"
    assert reloaded_ds.num_rows == ds.num_rows

    reloaded_mlflow_ds = mlflow.data.from_huggingface(
        reloaded_ds, path="cornell-movie-review-data/rotten_tomatoes"
    )
    assert reloaded_mlflow_ds.digest == mlflow_ds.digest


def test_from_huggingface_dataset_constructs_expected_dataset_with_revision():
    # Load this revision:
    # https://huggingface.co/datasets/cornell-movie-review-data/rotten_tomatoes/commit/aa13bc287fa6fcab6daf52f0dfb9994269ffea28
    revision = "aa13bc287fa6fcab6daf52f0dfb9994269ffea28"
    ds = datasets.load_dataset(
        "cornell-movie-review-data/rotten_tomatoes",
        split="train",
        revision=revision,
        trust_remote_code=True,
    )

    mlflow_ds_new = mlflow.data.from_huggingface(
        ds,
        path="cornell-movie-review-data/rotten_tomatoes",
        revision=revision,
        trust_remote_code=True,
    )

    mlflow_ds_new.source.load()
    assert mlflow_ds_new.source.revision == revision


def test_from_huggingface_dataset_constructs_expected_dataset_with_data_files():
    data_files = {"train": "prompts.csv"}
    ds = datasets.load_dataset("fka/awesome-chatgpt-prompts", data_files=data_files, split="train")
    mlflow_ds = mlflow.data.from_huggingface(
        ds, path="fka/awesome-chatgpt-prompts", data_files=data_files
    )

    assert isinstance(mlflow_ds, HuggingFaceDataset)
    assert mlflow_ds.ds == ds
    assert mlflow_ds.schema == _infer_schema(ds.to_pandas())
    assert mlflow_ds.profile == {
        "num_rows": ds.num_rows,
        "dataset_size": ds.dataset_size,
        "size_in_bytes": ds.size_in_bytes,
    }

    assert isinstance(mlflow_ds.source, HuggingFaceDatasetSource)
    reloaded_ds = mlflow_ds.source.load()
    assert reloaded_ds.builder_name == ds.builder_name
    assert reloaded_ds.config_name == ds.config_name
    assert reloaded_ds.split == ds.split == "train"
    assert reloaded_ds.num_rows == ds.num_rows

    reloaded_mlflow_ds = mlflow.data.from_huggingface(
        reloaded_ds, path="fka/awesome-chatgpt-prompts", data_files=data_files
    )
    assert reloaded_mlflow_ds.digest == mlflow_ds.digest


def test_from_huggingface_dataset_constructs_expected_dataset_with_data_dir(tmp_path):
    df = pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4, 5, 6]})
    data_dir = "data"
    os.makedirs(tmp_path / data_dir)
    df.to_csv(tmp_path / data_dir / "my_data.csv")
    ds = datasets.load_dataset(str(tmp_path), data_dir=data_dir, name="default", split="train")
    mlflow_ds = mlflow.data.from_huggingface(ds, path=str(tmp_path), data_dir=data_dir)

    assert mlflow_ds.ds == ds
    assert mlflow_ds.schema == _infer_schema(ds.to_pandas())
    assert mlflow_ds.profile == {
        "num_rows": ds.num_rows,
        "dataset_size": ds.dataset_size,
        "size_in_bytes": ds.size_in_bytes,
    }

    assert isinstance(mlflow_ds.source, HuggingFaceDatasetSource)
    reloaded_ds = mlflow_ds.source.load()
    assert reloaded_ds.builder_name == ds.builder_name
    assert reloaded_ds.config_name == ds.config_name
    assert reloaded_ds.split == ds.split == "train"
    assert reloaded_ds.num_rows == ds.num_rows

    reloaded_mlflow_ds = mlflow.data.from_huggingface(
        reloaded_ds, path=str(tmp_path), data_dir=data_dir
    )
    assert reloaded_mlflow_ds.digest == mlflow_ds.digest


def test_from_huggingface_dataset_respects_user_specified_name_and_digest():
    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
    mlflow_ds = mlflow.data.from_huggingface(
        ds, path="cornell-movie-review-data/rotten_tomatoes", name="myname", digest="mydigest"
    )
    assert mlflow_ds.name == "myname"
    assert mlflow_ds.digest == "mydigest"


def test_from_huggingface_dataset_digest_is_consistent_for_large_ordered_datasets(tmp_path):
    assert (
        mlflow.data.huggingface_dataset._MAX_ROWS_FOR_DIGEST_COMPUTATION_AND_SCHEMA_INFERENCE
        < 200000
    )

    df = pd.DataFrame.from_dict({
        "a": list(range(200000)),
        "b": list(range(200000)),
    })
    data_dir = "data"
    os.makedirs(tmp_path / data_dir)
    df.to_csv(tmp_path / data_dir / "my_data.csv")

    ds = datasets.load_dataset(str(tmp_path), data_dir=data_dir, name="default", split="train")
    mlflow_ds = mlflow.data.from_huggingface(ds, path=str(tmp_path), data_dir=data_dir)
    assert mlflow_ds.digest == "1dda4ce8"


def test_from_huggingface_dataset_throws_for_dataset_dict():
    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes")
    assert isinstance(ds, datasets.DatasetDict)

    with pytest.raises(
        MlflowException, match="must be an instance of `datasets.Dataset`.*DatasetDict"
    ):
        mlflow.data.from_huggingface(ds, path="cornell-movie-review-data/rotten_tomatoes")


def test_from_huggingface_dataset_no_source_specified():
    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
    mlflow_ds = mlflow.data.from_huggingface(ds)

    assert isinstance(mlflow_ds, HuggingFaceDataset)

    assert isinstance(mlflow_ds.source, CodeDatasetSource)
    assert "mlflow.source.name" in mlflow_ds.source.to_json()


def test_dataset_conversion_to_json():
    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
    mlflow_ds = mlflow.data.from_huggingface(ds, path="cornell-movie-review-data/rotten_tomatoes")

    dataset_json = mlflow_ds.to_json()
    parsed_json = json.loads(dataset_json)
    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
    assert parsed_json["name"] == mlflow_ds.name
    assert parsed_json["digest"] == mlflow_ds.digest
    assert parsed_json["source"] == mlflow_ds.source.to_json()
    assert parsed_json["source_type"] == mlflow_ds.source._get_source_type()
    assert parsed_json["profile"] == json.dumps(mlflow_ds.profile)

    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
    assert Schema.from_json(schema_json) == mlflow_ds.schema


def test_dataset_source_conversion_to_json():
    ds = datasets.load_dataset(
        "cornell-movie-review-data/rotten_tomatoes",
        split="train",
        revision="c33cbf965006dba64f134f7bef69c53d5d0d285d",
    )
    mlflow_ds = mlflow.data.from_huggingface(
        ds,
        path="cornell-movie-review-data/rotten_tomatoes",
        revision="c33cbf965006dba64f134f7bef69c53d5d0d285d",
    )
    source = mlflow_ds.source

    source_json = source.to_json()
    parsed_source = json.loads(source_json)
    assert parsed_source["revision"] == "c33cbf965006dba64f134f7bef69c53d5d0d285d"
    assert parsed_source["split"] == "train"
    assert parsed_source["config_name"] == "default"
    assert parsed_source["path"] == "cornell-movie-review-data/rotten_tomatoes"
    assert not parsed_source["data_dir"]
    assert not parsed_source["data_files"]

    reloaded_source = HuggingFaceDatasetSource.from_json(source_json)
    assert json.loads(reloaded_source.to_json()) == parsed_source

    reloaded_source = get_dataset_source_from_json(
        source_json, source_type=source._get_source_type()
    )
    assert isinstance(reloaded_source, HuggingFaceDatasetSource)
    assert type(source) == type(reloaded_source)
    assert reloaded_source.to_json() == source.to_json()


def test_to_evaluation_dataset():
    import numpy as np

    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
    dataset = mlflow.data.from_huggingface(
        ds, path="cornell-movie-review-data/rotten_tomatoes", targets="label"
    )

    evaluation_dataset = dataset.to_evaluation_dataset()
    assert isinstance(evaluation_dataset, EvaluationDataset)
    assert evaluation_dataset.features_data.equals(dataset.ds.to_pandas().drop("label", axis=1))
    assert np.array_equal(evaluation_dataset.labels_data, dataset.ds.to_pandas()["label"].values)


def test_from_huggingface_dataset_with_sample_source():
    source_uri = "test:/my/test/uri"
    source = SampleDatasetSource._resolve(source_uri)

    data = {"text": ["This is a sample text.", "Another sample text."], "label": [0, 1]}
    dataset = datasets.Dataset.from_dict(data)

    train_dataset = mlflow.data.from_huggingface(
        dataset,
        name="sample-text-dataset",
        source=source,
    )

    assert isinstance(train_dataset, HuggingFaceDataset)
    assert train_dataset.source == source