import json
from typing import Any
from unittest.mock import Mock

import pandas as pd
import pytest

from mlflow.data.dataset_source_registry import (
    get_dataset_source_from_json,
    register_dataset_source,
)
from mlflow.data.spark_dataset_source import SparkDatasetSource
from mlflow.entities.evaluation_dataset import DatasetGranularity
from mlflow.entities.evaluation_dataset import EvaluationDataset as MLflowEvaluationDataset
from mlflow.genai.datasets.databricks_evaluation_dataset_source import (
    DatabricksEvaluationDatasetSource,
    DatabricksUCTableDatasetSource,
)
from mlflow.genai.datasets.evaluation_dataset import EvaluationDataset


def create_test_source_json(table_name: str = "main.default.testtable") -> str:
    """Create a JSON string source value consistent with Databricks managed evaluation datasets.

    This format matches the behavior of Databricks managed evaluation datasets as of July 2025.
    """
    return json.dumps({"table_name": table_name})


def create_mock_managed_dataset(source_value: Any) -> Mock:
    """Create a mock Databricks Agent Evaluation ManagedDataset for testing"""
    mock_dataset = Mock()
    mock_dataset.dataset_id = getattr(source_value, "dataset_id", "test-dataset-id")
    mock_dataset.name = getattr(source_value, "_table_name", "catalog.schema.table")
    mock_dataset.digest = "test-digest"
    mock_dataset.schema = "test-schema"
    mock_dataset.profile = "test-profile"
    mock_dataset.source = source_value
    mock_dataset.source_type = "databricks-uc-table"
    mock_dataset.create_time = "2024-01-01T00:00:00"
    mock_dataset.created_by = "test-user"
    mock_dataset.last_update_time = "2024-01-02T00:00:00"
    mock_dataset.last_updated_by = "test-user-2"

    # Mock methods
    mock_dataset.to_df.return_value = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    mock_dataset.set_profile.return_value = mock_dataset
    mock_dataset.merge_records.return_value = mock_dataset

    return mock_dataset


@pytest.fixture
def mock_managed_dataset() -> Mock:
    """Create a mock Databricks Agent Evaluation ManagedDataset for testing."""
    return create_mock_managed_dataset(create_test_source_json())


def create_dataset_with_source(source_value: Any) -> EvaluationDataset:
    """Factory function to create EvaluationDataset with specific source value."""
    mock_dataset = create_mock_managed_dataset(source_value)
    return EvaluationDataset(mock_dataset)


def test_evaluation_dataset_properties(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    assert dataset.dataset_id == "test-dataset-id"
    assert dataset.name == "catalog.schema.table"
    assert dataset.digest == "test-digest"
    assert dataset.schema == "test-schema"
    assert dataset.profile == "test-profile"
    assert dataset.source_type == "databricks-uc-table"
    assert dataset.create_time == "2024-01-01T00:00:00"
    assert dataset.created_by == "test-user"
    assert dataset.last_update_time == "2024-01-02T00:00:00"
    assert dataset.last_updated_by == "test-user-2"
    assert isinstance(dataset.source, DatabricksEvaluationDatasetSource)
    assert dataset.source.table_name == "catalog.schema.table"
    assert dataset.source.dataset_id == "test-dataset-id"


def test_evaluation_dataset_source_with_string_source():
    dataset = create_dataset_with_source("string-value")

    assert isinstance(dataset.source, DatabricksEvaluationDatasetSource)
    assert dataset.source.table_name == "catalog.schema.table"
    assert dataset.source.dataset_id == "test-dataset-id"


def test_evaluation_dataset_source_with_none():
    dataset = create_dataset_with_source(None)

    assert isinstance(dataset.source, DatabricksEvaluationDatasetSource)
    assert dataset.source.table_name == "catalog.schema.table"
    assert dataset.source.dataset_id == "test-dataset-id"


def test_evaluation_dataset_source_always_returns_databricks_evaluation_dataset_source():
    existing_source = DatabricksEvaluationDatasetSource(
        table_name="existing.table", dataset_id="existing-id"
    )
    dataset = create_dataset_with_source(existing_source)

    assert isinstance(dataset.source, DatabricksEvaluationDatasetSource)
    assert dataset.source.table_name == "existing.table"
    assert dataset.source.dataset_id == "existing-id"

    spark_source = SparkDatasetSource(table_name="spark.table")
    dataset = create_dataset_with_source(spark_source)

    assert isinstance(dataset.source, DatabricksEvaluationDatasetSource)
    assert dataset.source.table_name == "spark.table"
    assert dataset.source.dataset_id == "test-dataset-id"


def test_evaluation_dataset_to_df(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    df = dataset.to_df()
    assert isinstance(df, pd.DataFrame)
    assert len(df) == 3
    mock_managed_dataset.to_df.assert_called_once()


def test_evaluation_dataset_to_mlflow_entity(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    entity = dataset._to_mlflow_entity()
    assert entity.name == "catalog.schema.table"
    assert entity.digest == "test-digest"
    assert entity.source_type == "databricks-uc-table"

    source_dict = json.loads(entity.source)
    assert source_dict["table_name"] == "catalog.schema.table"
    assert source_dict["dataset_id"] == "test-dataset-id"
    assert entity.schema == "test-schema"
    assert entity.profile == "test-profile"


def test_evaluation_dataset_to_mlflow_entity_with_existing_source():
    existing_source = DatabricksEvaluationDatasetSource(
        table_name="existing.table", dataset_id="existing-id"
    )
    dataset = create_dataset_with_source(existing_source)

    entity = dataset._to_mlflow_entity()
    assert entity.name == "existing.table"
    assert entity.digest == "test-digest"
    assert entity.source_type == "databricks-uc-table"

    source_dict = json.loads(entity.source)
    assert source_dict["table_name"] == "existing.table"
    assert source_dict["dataset_id"] == "existing-id"
    assert entity.schema == "test-schema"
    assert entity.profile == "test-profile"


def test_evaluation_dataset_set_profile(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    new_dataset = dataset.set_profile("new-profile")
    assert isinstance(new_dataset, EvaluationDataset)
    mock_managed_dataset.set_profile.assert_called_once_with("new-profile")


def test_evaluation_dataset_merge_records(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    new_records = [{"col1": 4, "col2": "d"}]
    new_dataset = dataset.merge_records(new_records)
    assert isinstance(new_dataset, EvaluationDataset)
    mock_managed_dataset.merge_records.assert_called_once_with(new_records)


def test_evaluation_dataset_delete_records(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    record_ids = ["record-1", "record-2"]
    dataset.delete_records(record_ids)
    mock_managed_dataset.delete_records.assert_called_once_with(record_ids)


def test_evaluation_dataset_digest_computation(mock_managed_dataset):
    # Test when managed dataset has no digest
    mock_managed_dataset.digest = None

    dataset = EvaluationDataset(mock_managed_dataset)
    digest = dataset.digest

    assert digest is not None


def test_evaluation_dataset_to_evaluation_dataset(mock_managed_dataset):
    dataset = EvaluationDataset(mock_managed_dataset)

    legacy_dataset = dataset.to_evaluation_dataset(
        path="/path/to/data", feature_names=["col1", "col2"]
    )

    assert legacy_dataset._features_data.equals(dataset.to_df())
    assert legacy_dataset._path == "/path/to/data"
    assert legacy_dataset._feature_names == ["col1", "col2"]
    assert legacy_dataset.name == "catalog.schema.table"
    assert legacy_dataset.digest == "test-digest"


def test_databricks_uc_table_dataset_source():
    register_dataset_source(DatabricksUCTableDatasetSource)

    source_json = json.dumps({"table_name": "catalog.schema.table", "dataset_id": "test-id"})

    source = get_dataset_source_from_json(source_json, "databricks-uc-table")
    assert isinstance(source, DatabricksUCTableDatasetSource)
    assert source._get_source_type() == "databricks-uc-table"
    assert source.table_name == "catalog.schema.table"
    assert source.dataset_id == "test-id"


def _create_mlflow_evaluation_dataset() -> MLflowEvaluationDataset:
    return MLflowEvaluationDataset(
        dataset_id="test-id",
        name="test-dataset",
        digest="test-digest",
        created_time=0,
        last_update_time=0,
    )


@pytest.mark.parametrize(
    ("input_keys", "expected_granularity"),
    [
        # empty keys -> UNKNOWN
        (set(), DatasetGranularity.UNKNOWN),
        # no 'goal' field -> TRACE
        ({"request"}, DatasetGranularity.TRACE),
        ({"messages"}, DatasetGranularity.TRACE),
        ({"query", "context"}, DatasetGranularity.TRACE),
        # 'goal' and only session fields -> SESSION
        ({"goal"}, DatasetGranularity.SESSION),
        ({"goal", "persona"}, DatasetGranularity.SESSION),
        ({"goal", "context"}, DatasetGranularity.SESSION),
        ({"goal", "persona", "context"}, DatasetGranularity.SESSION),
        # 'goal' mixed with non-session fields -> UNKNOWN
        ({"goal", "request"}, DatasetGranularity.UNKNOWN),
        ({"goal", "messages"}, DatasetGranularity.UNKNOWN),
        ({"goal", "persona", "extra_field"}, DatasetGranularity.UNKNOWN),
    ],
)
def test_classify_input_fields(
    input_keys,
    expected_granularity,
):
    dataset = _create_mlflow_evaluation_dataset()
    result = dataset._classify_input_fields(input_keys)
    assert result == expected_granularity