from textwrap import dedent
from unittest.mock import MagicMock
from uuid import uuid4

import pytest

from .classification_model import ClassificationMetrics, ClassificationModel
from .datasource import Datasource
from .embedding_model import PretrainedEmbeddingModel
from .memoryset import LabeledMemoryset
from .telemetry import ClassificationPrediction


def test_create_model(classification_model: ClassificationModel, readonly_memoryset: LabeledMemoryset):
    assert classification_model is not None
    assert classification_model.name == "test_classification_model"
    assert classification_model.memoryset == readonly_memoryset
    assert classification_model.num_classes == 2
    assert classification_model.memory_lookup_count == 3


def test_create_model_already_exists_error(readonly_memoryset, classification_model):
    with pytest.raises(ValueError):
        ClassificationModel.create("test_classification_model", readonly_memoryset)
    with pytest.raises(ValueError):
        ClassificationModel.create("test_classification_model", readonly_memoryset, if_exists="error")


def test_create_model_already_exists_return(readonly_memoryset, classification_model):
    with pytest.raises(ValueError):
        ClassificationModel.create(
            "test_classification_model", readonly_memoryset, if_exists="open", balance_classes=True
        )

    with pytest.raises(ValueError):
        ClassificationModel.create(
            "test_classification_model", readonly_memoryset, if_exists="open", memory_lookup_count=37
        )

    with pytest.raises(ValueError):
        ClassificationModel.create("test_classification_model", readonly_memoryset, if_exists="open", num_classes=19)

    with pytest.raises(ValueError):
        ClassificationModel.create(
            "test_classification_model", readonly_memoryset, if_exists="open", min_memory_weight=0.77
        )

    new_model = ClassificationModel.create("test_classification_model", readonly_memoryset, if_exists="open")
    assert new_model is not None
    assert new_model.name == "test_classification_model"
    assert new_model.memoryset == readonly_memoryset
    assert new_model.num_classes == 2
    assert new_model.memory_lookup_count == 3


def test_create_model_unauthenticated(unauthenticated_client, readonly_memoryset: LabeledMemoryset):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            ClassificationModel.create("test_model", readonly_memoryset)


def test_get_model(classification_model: ClassificationModel):
    fetched_model = ClassificationModel.open(classification_model.name)
    assert fetched_model is not None
    assert fetched_model.id == classification_model.id
    assert fetched_model.name == classification_model.name
    assert fetched_model.num_classes == 2
    assert fetched_model.memory_lookup_count == 3
    assert fetched_model == classification_model


def test_get_model_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            ClassificationModel.open("test_model")


def test_get_model_invalid_input():
    with pytest.raises(ValueError, match="Invalid input"):
        ClassificationModel.open("not valid id")


def test_get_model_not_found():
    with pytest.raises(LookupError):
        ClassificationModel.open(str(uuid4()))


def test_get_model_unauthorized(unauthorized_client, classification_model: ClassificationModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            ClassificationModel.open(classification_model.name)


def test_list_models(classification_model: ClassificationModel):
    models = ClassificationModel.all()
    assert len(models) > 0
    assert any(model.name == model.name for model in models)


def test_list_models_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            ClassificationModel.all()


def test_list_models_unauthorized(unauthorized_client, classification_model: ClassificationModel):
    with unauthorized_client.use():
        assert ClassificationModel.all() == []


def test_memoryset_classification_models_property(
    classification_model: ClassificationModel, readonly_memoryset: LabeledMemoryset
):
    models = readonly_memoryset.classification_models
    assert len(models) > 0
    assert any(model.id == classification_model.id for model in models)


def test_update_model_attributes(classification_model: ClassificationModel):
    original_description = classification_model.description

    classification_model.set(description=None)
    assert classification_model.description is None

    classification_model.set(description=original_description)
    assert classification_model.description == original_description

    classification_model.set(locked=True)
    assert classification_model.locked is True

    classification_model.set(locked=False)
    assert classification_model.locked is False

    classification_model.lock()
    assert classification_model.locked is True

    classification_model.unlock()
    assert classification_model.locked is False


def test_delete_model(readonly_memoryset: LabeledMemoryset):
    ClassificationModel.create("model_to_delete", LabeledMemoryset.open(readonly_memoryset.name))
    assert ClassificationModel.open("model_to_delete")
    ClassificationModel.drop("model_to_delete")
    with pytest.raises(LookupError):
        ClassificationModel.open("model_to_delete")


def test_delete_model_unauthenticated(unauthenticated_client, classification_model: ClassificationModel):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            ClassificationModel.drop(classification_model.name)


def test_delete_model_not_found():
    with pytest.raises(LookupError):
        ClassificationModel.drop(str(uuid4()))
    # ignores error if specified
    ClassificationModel.drop(str(uuid4()), if_not_exists="ignore")


def test_delete_model_unauthorized(unauthorized_client, classification_model: ClassificationModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            ClassificationModel.drop(classification_model.name)


def test_delete_memoryset_before_model_constraint_violation(datasource):
    memoryset = LabeledMemoryset.from_datasource(
        "test_memoryset_delete_before_model", datasource=datasource, label_names=["soup", "cats"]
    )
    ClassificationModel.create("test_model_delete_before_memoryset", memoryset)
    with pytest.raises(RuntimeError):
        LabeledMemoryset.drop(memoryset.id)


def test_delete_memoryset_with_model_cascade(datasource):
    """Test that cascade=False prevents deletion and cascade=True allows it."""
    memoryset = LabeledMemoryset.from_datasource(
        "test_memoryset_cascade_delete", datasource=datasource, label_names=["soup", "cats"]
    )
    model = ClassificationModel.create("test_model_cascade_delete", memoryset)

    # Verify model exists
    assert ClassificationModel.open(model.name) is not None

    # Without cascade, deletion should fail
    with pytest.raises(RuntimeError):
        LabeledMemoryset.drop(memoryset.id, cascade=False)

    # Model should still exist
    assert ClassificationModel.exists(model.name)

    # With cascade, deletion should succeed
    LabeledMemoryset.drop(memoryset.id, cascade=True)

    # Model should be deleted along with the memoryset
    assert not ClassificationModel.exists(model.name)
    assert not LabeledMemoryset.exists(memoryset.name)


def test_evaluate(classification_model, eval_datasource: Datasource):
    result = classification_model.evaluate(eval_datasource)

    assert result is not None
    assert isinstance(result, ClassificationMetrics)

    assert isinstance(result.accuracy, float)
    assert result.accuracy == pytest.approx(0.5)
    assert isinstance(result.f1_score, float)
    assert result.f1_score == pytest.approx(0.5)
    assert isinstance(result.loss, float)

    assert isinstance(result.anomaly_score_mean, float)
    assert isinstance(result.anomaly_score_median, float)
    assert isinstance(result.anomaly_score_variance, float)
    assert -1.0 <= result.anomaly_score_mean <= 1.0
    assert -1.0 <= result.anomaly_score_median <= 1.0
    assert -1.0 <= result.anomaly_score_variance <= 1.0

    assert result.pr_auc is not None
    assert result.pr_auc == pytest.approx(0.83333, rel=1e-5)
    assert result.pr_curve is not None
    assert result.pr_curve["thresholds"] == pytest.approx(
        [0.0, 0.3021204173564911, 0.30852025747299194, 0.6932827234268188, 0.6972201466560364]
    )
    assert result.pr_curve["precisions"] == pytest.approx([0.5, 2 / 3, 0.5, 1.0, 1.0])
    assert result.pr_curve["recalls"] == pytest.approx([1.0, 1.0, 0.5, 0.5, 0.0])

    assert result.roc_auc is not None
    assert result.roc_auc == pytest.approx(0.75)
    assert result.roc_curve is not None
    assert result.roc_curve["thresholds"] == pytest.approx(
        [0.3021204173564911, 0.30852025747299194, 0.6932827234268188, 0.6972201466560364, 1.0]
    )
    assert result.roc_curve["false_positive_rates"] == pytest.approx([1.0, 0.5, 0.5, 0.0, 0.0])
    assert result.roc_curve["true_positive_rates"] == pytest.approx([1.0, 1.0, 0.5, 0.5, 0.0])


def test_evaluate_local_dataset(classification_model, eval_data: list[dict]):
    datasets = pytest.importorskip("datasets")
    eval_dataset = datasets.Dataset.from_list(eval_data)
    result = classification_model.evaluate(eval_dataset)
    assert result is not None
    assert isinstance(result, ClassificationMetrics)
    assert isinstance(result.accuracy, float)
    assert isinstance(result.f1_score, float)
    assert isinstance(result.loss, float)


def test_evaluate_local_list(classification_model, eval_data: list[dict]):
    result = classification_model.evaluate(eval_data)
    assert result is not None
    assert isinstance(result, ClassificationMetrics)
    assert isinstance(result.accuracy, float)
    assert isinstance(result.f1_score, float)
    assert isinstance(result.loss, float)


def test_evaluate_datasource_with_nones_raises_error(classification_model: ClassificationModel, datasource: Datasource):
    with pytest.raises(ValueError):
        classification_model.evaluate(datasource, record_predictions=True, tags={"test"})


def test_evaluate_with_telemetry(classification_model: ClassificationModel, eval_data: list[dict]):
    result = classification_model.evaluate(eval_data, record_predictions=True, tags={"test"}, batch_size=2)
    assert result is not None
    assert isinstance(result, ClassificationMetrics)
    predictions = classification_model.predictions(tag="test", batch_size=100, sort=[("timestamp", "asc")])
    assert len(predictions) == 4
    assert all(p.tags == {"test"} for p in predictions)
    prediction_expected_labels = [p.expected_label if p.expected_label is not None else -1 for p in predictions]
    eval_expected_labels = [d["label"] for d in eval_data]
    assert all(
        p == l for p, l in zip(prediction_expected_labels, eval_expected_labels)
    ), f"Prediction expected labels: {prediction_expected_labels} do not match eval expected labels: {eval_expected_labels}"


def test_evaluate_with_partition_column_include_global(partitioned_classification_model: ClassificationModel):
    result = partitioned_classification_model.evaluate(
        [
            {"value": "soup is good", "label": 0, "partition_id": "p1"},
            {"value": "cats are cute", "label": 1, "partition_id": "p1"},
        ],
        partition_column="partition_id",
        partition_filter_mode="include_global",
    )
    assert result is not None
    assert isinstance(result, ClassificationMetrics)


def test_evaluate_with_partition_column_exclude_global(partitioned_classification_model: ClassificationModel):
    result = partitioned_classification_model.evaluate(
        [
            {"value": "soup is good", "label": 0, "partition_id": "p1"},
            {"value": "cats are cute", "label": 1, "partition_id": "p1"},
        ],
        partition_column="partition_id",
        partition_filter_mode="exclude_global",
    )
    assert result is not None
    assert isinstance(result, ClassificationMetrics)


def test_evaluate_with_partition_column_only_global(partitioned_classification_model: ClassificationModel):
    result = partitioned_classification_model.evaluate(
        [
            {"value": "cats are independent animals", "label": 1, "partition_id": None},
            {"value": "i love the beach", "label": 1, "partition_id": None},
        ],
        partition_column="partition_id",
        partition_filter_mode="only_global",
    )
    assert result is not None
    assert isinstance(result, ClassificationMetrics)


def test_evaluate_with_partition_column_ignore_partitions(partitioned_classification_model: ClassificationModel):
    result = partitioned_classification_model.evaluate(
        [
            {"value": "soup is good", "label": 0, "partition_id": "p1"},
            {"value": "cats are cute", "label": 1, "partition_id": "p2"},
        ],
        partition_column="partition_id",
        partition_filter_mode="ignore_partitions",
    )
    assert result is not None
    assert isinstance(result, ClassificationMetrics)


def test_evaluate_with_partition_column_datasource(partitioned_classification_model: ClassificationModel):
    eval_datasource = Datasource.from_list(
        "eval_partition_datasource",
        [
            {"value": "soup is good", "label": 0, "partition_id": "p1"},
            {"value": "cats are cute", "label": 1, "partition_id": "p1"},
            {"value": "homemade soup recipes", "label": 0, "partition_id": "p2"},
            {"value": "cats purr when happy", "label": 1, "partition_id": "p2"},
        ],
    )
    result = partitioned_classification_model.evaluate(
        eval_datasource,
        partition_column="partition_id",
        partition_filter_mode="exclude_global",
    )
    assert result is not None
    assert isinstance(result, ClassificationMetrics)


def test_predict(classification_model: ClassificationModel, label_names: list[str]):
    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"], batch_size=1)
    assert len(predictions) == 2
    assert predictions[0].prediction_id is not None
    assert predictions[1].prediction_id is not None
    assert predictions[0].label == 0
    assert predictions[0].label_name == label_names[0]
    assert 0 <= predictions[0].confidence <= 1
    assert predictions[1].label == 1
    assert predictions[1].label_name == label_names[1]
    assert 0 <= predictions[1].confidence <= 1

    assert predictions[0].logits is not None
    assert predictions[1].logits is not None
    assert len(predictions[0].logits) == 2
    assert len(predictions[1].logits) == 2
    assert predictions[0].logits[0] > predictions[0].logits[1]
    assert predictions[1].logits[0] < predictions[1].logits[1]


def test_classification_prediction_has_no_label(classification_model: ClassificationModel):
    """Ensure optional score is None for classification predictions."""
    prediction = classification_model.predict("Do you want to go to the beach?")
    assert isinstance(prediction, ClassificationPrediction)
    assert prediction.label is None


def test_predict_disable_telemetry(classification_model: ClassificationModel, label_names: list[str]):
    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"], save_telemetry="off")
    assert len(predictions) == 2
    assert predictions[0].prediction_id is None
    assert predictions[1].prediction_id is None
    assert predictions[0].label == 0
    assert predictions[0].label_name == label_names[0]
    assert 0 <= predictions[0].confidence <= 1
    assert predictions[1].label == 1
    assert predictions[1].label_name == label_names[1]
    assert 0 <= predictions[1].confidence <= 1


def test_predict_unauthenticated(unauthenticated_client, classification_model: ClassificationModel):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            classification_model.predict(["Do you love soup?", "Are cats cute?"])


def test_predict_unauthorized(unauthorized_client, classification_model: ClassificationModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            classification_model.predict(["Do you love soup?", "Are cats cute?"])


def test_predict_constraint_violation(readonly_memoryset: LabeledMemoryset):
    model = ClassificationModel.create(
        "test_model_lookup_count_too_high",
        readonly_memoryset,
        num_classes=2,
        memory_lookup_count=readonly_memoryset.length + 2,
    )
    with pytest.raises(RuntimeError):
        model.predict("test")


def test_predict_with_partition_id(partitioned_classification_model: ClassificationModel, label_names: list[str]):
    """Test predict with a specific partition_id"""
    # Predict with partition_id p1 - should use memories from p1
    prediction = partitioned_classification_model.predict(
        "soup", partition_id="p1", partition_filter_mode="exclude_global"
    )
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1
    assert prediction.logits is not None
    assert len(prediction.logits) == 2

    # Predict with partition_id p2 - should use memories from p2
    prediction_p2 = partitioned_classification_model.predict(
        "cats", partition_id="p2", partition_filter_mode="exclude_global"
    )
    assert prediction_p2.label is not None
    assert prediction_p2.label_name in label_names
    assert 0 <= prediction_p2.confidence <= 1


def test_predict_with_partition_id_include_global(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test predict with partition_id and include_global mode (default)"""
    # Predict with partition_id p1 and include_global (default) - should include both p1 and global memories
    prediction = partitioned_classification_model.predict(
        "soup", partition_id="p1", partition_filter_mode="include_global"
    )
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1


def test_predict_with_partition_id_exclude_global(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test predict with partition_id and exclude_global mode"""
    # Predict with partition_id p1 and exclude_global - should only use p1 memories
    prediction = partitioned_classification_model.predict(
        "soup", partition_id="p1", partition_filter_mode="exclude_global"
    )
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1


def test_predict_with_partition_id_only_global(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test predict with partition_filter_mode only_global"""
    # Predict with only_global mode - should only use global memories
    prediction = partitioned_classification_model.predict("cats", partition_filter_mode="only_global")
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1


def test_predict_with_partition_id_ignore_partitions(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test predict with partition_filter_mode ignore_partitions"""
    # Predict with ignore_partitions mode - should ignore partition filtering
    prediction = partitioned_classification_model.predict("soup", partition_filter_mode="ignore_partitions")
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1


def test_predict_batch_with_partition_id(partitioned_classification_model: ClassificationModel, label_names: list[str]):
    """Test batch predict with partition_id"""
    # Batch predict with partition_id p1
    predictions = partitioned_classification_model.predict(
        ["soup is good", "cats are cute"],
        partition_id="p1",
        partition_filter_mode="exclude_global",
    )
    assert len(predictions) == 2
    assert all(p.label is not None for p in predictions)
    assert all(p.label_name in label_names for p in predictions)
    assert all(0 <= p.confidence <= 1 for p in predictions)
    assert all(p.logits is not None and len(p.logits) == 2 for p in predictions)


def test_predict_with_partition_id_and_filters(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test predict with partition_id and filters"""
    # Predict with partition_id and filters
    prediction = partitioned_classification_model.predict(
        "soup",
        partition_id="p1",
        partition_filter_mode="exclude_global",
        filters=[("key", "==", "g1")],
    )
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1


def test_predict_batch_with_list_of_partition_ids(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test batch predict with a list of partition_ids (one for each query input)"""
    # Batch predict with a list of partition_ids - one for each input
    # First input uses p1, second input uses p2
    predictions = partitioned_classification_model.predict(
        ["soup is good", "cats are cute"],
        partition_id=["p1", "p2"],
        partition_filter_mode="exclude_global",
    )
    assert len(predictions) == 2
    assert all(p.label is not None for p in predictions)
    assert all(p.label_name in label_names for p in predictions)
    assert all(0 <= p.confidence <= 1 for p in predictions)
    assert all(p.logits is not None and len(p.logits) == 2 for p in predictions)

    # Verify that predictions were made using the correct partitions
    # Each prediction should use memories from its respective partition
    assert predictions[0].input_value == "soup is good"
    assert predictions[1].input_value == "cats are cute"


@pytest.mark.asyncio
async def test_predict_async_with_partition_id(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test async predict with partition_id"""
    # Async predict with partition_id p1
    prediction = await partitioned_classification_model.apredict(
        "soup", partition_id="p1", partition_filter_mode="exclude_global"
    )
    assert prediction.label is not None
    assert prediction.label_name in label_names
    assert 0 <= prediction.confidence <= 1
    assert prediction.logits is not None
    assert len(prediction.logits) == 2


@pytest.mark.asyncio
async def test_predict_async_batch_with_partition_id(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test async batch predict with partition_id"""
    # Async batch predict with partition_id p1
    predictions = await partitioned_classification_model.apredict(
        ["soup is good", "cats are cute"],
        partition_id="p1",
        partition_filter_mode="exclude_global",
    )
    assert len(predictions) == 2
    assert all(p.label is not None for p in predictions)
    assert all(p.label_name in label_names for p in predictions)
    assert all(0 <= p.confidence <= 1 for p in predictions)


@pytest.mark.asyncio
async def test_predict_async_batch_with_list_of_partition_ids(
    partitioned_classification_model: ClassificationModel, label_names: list[str]
):
    """Test async batch predict with a list of partition_ids (one for each query input)"""
    # Async batch predict with a list of partition_ids - one for each input
    # First input uses p1, second input uses p2
    predictions = await partitioned_classification_model.apredict(
        ["soup is good", "cats are cute"],
        partition_id=["p1", "p2"],
        partition_filter_mode="exclude_global",
    )
    assert len(predictions) == 2
    assert all(p.label is not None for p in predictions)
    assert all(p.label_name in label_names for p in predictions)
    assert all(0 <= p.confidence <= 1 for p in predictions)
    assert all(p.logits is not None and len(p.logits) == 2 for p in predictions)

    # Verify that predictions were made using the correct partitions
    # Each prediction should use memories from its respective partition
    assert predictions[0].input_value == "soup is good"
    assert predictions[1].input_value == "cats are cute"


def test_record_prediction_feedback(classification_model: ClassificationModel):
    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"])
    expected_labels = [0, 1]
    classification_model.record_feedback(
        {
            "prediction_id": p.prediction_id,
            "category": "correct",
            "value": p.label == expected_label,
        }
        for expected_label, p in zip(expected_labels, predictions)
    )


def test_record_prediction_feedback_missing_category(classification_model: ClassificationModel):
    prediction = classification_model.predict("Do you love soup?")
    with pytest.raises(ValueError):
        classification_model.record_feedback({"prediction_id": prediction.prediction_id, "value": True})


def test_record_prediction_feedback_invalid_value(classification_model: ClassificationModel):
    prediction = classification_model.predict("Do you love soup?")
    with pytest.raises(ValueError, match=r"Invalid input.*"):
        classification_model.record_feedback(
            {"prediction_id": prediction.prediction_id, "category": "correct", "value": "invalid"}
        )


def test_record_prediction_feedback_invalid_prediction_id(classification_model: ClassificationModel):
    with pytest.raises(ValueError, match=r"Invalid input.*"):
        classification_model.record_feedback({"prediction_id": "invalid", "category": "correct", "value": True})


def test_predict_with_memoryset_override(
    classification_model: ClassificationModel, data: list[dict], label_names: list[str]
):
    inverted_labeled_memoryset = LabeledMemoryset.from_list(
        "test_memoryset_inverted_labels",
        [{"value": x["value"], "label": 1 if x["label"] == 0 else 0} for x in data],
        label_names=label_names,
        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
    )
    with classification_model.use_memoryset(inverted_labeled_memoryset):
        predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"], tags={"override_test"})
        assert predictions[0].label == 1
        assert predictions[1].label == 0
        assert predictions[0].memoryset == inverted_labeled_memoryset
        assert predictions[1].memoryset == inverted_labeled_memoryset

    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"], tags={"no_override_test"})
    assert predictions[0].label == 0
    assert predictions[1].label == 1
    assert predictions[0].memoryset == classification_model.memoryset
    assert predictions[1].memoryset == classification_model.memoryset

    # Verify historical predictions have the correct memoryset from telemetry
    override_predictions = classification_model.predictions(tag="override_test")
    assert len(override_predictions) == 2
    assert override_predictions[0].memoryset == inverted_labeled_memoryset
    assert override_predictions[1].memoryset == inverted_labeled_memoryset

    no_override_predictions = classification_model.predictions(tag="no_override_test")
    assert len(no_override_predictions) == 2
    assert no_override_predictions[0].memoryset == classification_model.memoryset
    assert no_override_predictions[1].memoryset == classification_model.memoryset


def test_predict_with_expected_labels(classification_model: ClassificationModel):
    prediction = classification_model.predict("Do you love soup?", expected_labels=1)
    assert prediction.expected_label == 1


def test_predict_with_expected_labels_no_telemetry(classification_model: ClassificationModel):
    """Test that expected_label is available even when telemetry is disabled"""
    prediction = classification_model.predict("Do you love soup?", expected_labels=1, save_telemetry="off")
    assert prediction.prediction_id is None  # telemetry is off
    assert prediction.expected_label == 1  # but expected_label should still be available


def test_predict_with_expected_labels_invalid_input(classification_model: ClassificationModel):
    # invalid number of expected labels for batch prediction
    with pytest.raises(ValueError, match=r"Invalid input.*"):
        classification_model.predict(["Do you love soup?", "Are cats cute?"], expected_labels=[0])
    # invalid label value
    with pytest.raises(ValueError):
        classification_model.predict("Do you love soup?", expected_labels=5)


def test_predict_with_filters(classification_model: ClassificationModel):
    # there are no memories with label 0 and key g1, so we force a wrong prediction
    filtered_prediction = classification_model.predict("I love soup", filters=[("key", "==", "g2")])
    assert filtered_prediction.label == 1
    assert filtered_prediction.label_name == "cats"


def test_predict_with_memoryset_update(writable_memoryset: LabeledMemoryset):
    model = ClassificationModel.create(
        "test_predict_with_memoryset_update",
        writable_memoryset,
        num_classes=2,
        memory_lookup_count=3,
    )
    try:
        prediction = model.predict("Do you love soup?", partition_filter_mode="ignore_partitions")
        assert prediction.label == 0
        assert prediction.label_name == "soup"
        # insert new memories
        writable_memoryset.insert(
            [
                {"value": "Do you love soup?", "label": 1, "key": "g1"},
                {"value": "Do you love soup for dinner?", "label": 1, "key": "g2"},
                {"value": "Do you love crackers?", "label": 1, "key": "g2"},
                {"value": "Do you love broth?", "label": 1, "key": "g2"},
                {"value": "Do you love chicken soup?", "label": 1, "key": "g2"},
                {"value": "Do you love chicken soup for dinner?", "label": 1, "key": "g2"},
                {"value": "Do you love chicken soup for dinner?", "label": 1, "key": "g2"},
            ],
        )
        prediction = model.predict("Do you love soup?")
        assert prediction.label == 1
        assert prediction.label_name == "cats"
    finally:
        ClassificationModel.drop("test_predict_with_memoryset_update")


def test_last_prediction_with_batch(classification_model: ClassificationModel):
    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"])
    assert classification_model.last_prediction is not None
    assert classification_model.last_prediction.prediction_id == predictions[-1].prediction_id
    assert classification_model.last_prediction.input_value == "Are cats cute?"
    assert classification_model._last_prediction_was_batch is True


def test_last_prediction_with_single(classification_model: ClassificationModel):
    # Test that last_prediction is updated correctly with single prediction
    prediction = classification_model.predict("Do you love soup?")
    assert classification_model.last_prediction is not None
    assert classification_model.last_prediction.prediction_id == prediction.prediction_id
    assert classification_model.last_prediction.input_value == "Do you love soup?"
    assert classification_model._last_prediction_was_batch is False


def test_explain(classification_model: ClassificationModel, requires_agentic_features):
    prediction = classification_model.predict("Do you love soup?")
    explanation = prediction.explanation
    assert explanation is not None
    assert len(explanation) > 10
    assert "soup" in explanation.lower()


def test_action_recommendation(
    classification_model: ClassificationModel, writable_memoryset: LabeledMemoryset, requires_agentic_features
):
    with classification_model.use_memoryset(writable_memoryset):
        prediction = classification_model.predict("Do you love soup?", expected_labels=1)
    assert prediction.memoryset == writable_memoryset

    memoryset_length = writable_memoryset.length

    # Get action recommendation
    action, rationale = prediction.recommend_action()

    assert action is not None
    assert rationale is not None
    assert action in ["remove_duplicates", "detect_mislabels", "add_memories", "finetuning"]
    assert len(rationale) > 10

    # Test memory suggestions
    suggestions = prediction.generate_memory_suggestions(num_memories=2)

    assert len(suggestions) == 2
    for suggestion in suggestions:
        assert isinstance(suggestion.value, str)
        assert len(suggestion.value) > 0
        assert isinstance(suggestion.label_name, str)
        assert suggestion.label_name in writable_memoryset.label_names

    writable_memoryset.insert(suggestions)

    writable_memoryset.refresh()
    assert writable_memoryset.length == memoryset_length + 2


def test_predict_with_prompt(classification_model: ClassificationModel):
    """Test that prompt parameter is properly passed through to predictions"""
    # Test with an instruction-supporting embedding model if available
    prediction_with_prompt = classification_model.predict(
        "I love this product!", prompt="Represent this text for sentiment classification:"
    )
    prediction_without_prompt = classification_model.predict("I love this product!")

    # Both should work and return valid predictions
    assert prediction_with_prompt.label is not None
    assert prediction_without_prompt.label is not None


def test_predict_with_empty_partition(fully_partitioned_classification_resources):
    datasource, memoryset, classification_model = fully_partitioned_classification_resources

    assert memoryset.length == 15

    with pytest.raises(RuntimeError, match="lookup failed to return the correct number of memories"):
        classification_model.predict("i love cats", partition_filter_mode="only_global")

    with pytest.raises(RuntimeError, match="lookup failed to return the correct number of memories"):
        classification_model.predict(
            "i love cats", partition_filter_mode="exclude_global", partition_id="p_does_not_exist"
        )

    with pytest.raises(RuntimeError, match="lookup failed to return the correct number of memories"):
        classification_model.evaluate(datasource, partition_filter_mode="only_global")


@pytest.mark.asyncio
async def test_predict_async_single(classification_model: ClassificationModel, label_names: list[str]):
    """Test async prediction with a single value"""
    prediction = await classification_model.apredict("Do you love soup?")
    assert isinstance(prediction, ClassificationPrediction)
    assert prediction.prediction_id is not None
    assert prediction.label == 0
    assert prediction.label_name == label_names[0]
    assert 0 <= prediction.confidence <= 1
    assert prediction.logits is not None
    assert len(prediction.logits) == 2


@pytest.mark.asyncio
async def test_predict_async_batch(classification_model: ClassificationModel, label_names: list[str]):
    """Test async prediction with a batch of values"""
    predictions = await classification_model.apredict(["Do you love soup?", "Are cats cute?"])
    assert len(predictions) == 2
    assert predictions[0].prediction_id is not None
    assert predictions[1].prediction_id is not None
    assert predictions[0].label == 0
    assert predictions[0].label_name == label_names[0]
    assert 0 <= predictions[0].confidence <= 1
    assert predictions[1].label == 1
    assert predictions[1].label_name == label_names[1]
    assert 0 <= predictions[1].confidence <= 1


@pytest.mark.asyncio
async def test_predict_async_with_expected_labels(classification_model: ClassificationModel):
    """Test async prediction with expected labels"""
    prediction = await classification_model.apredict("Do you love soup?", expected_labels=1)
    assert prediction.expected_label == 1


@pytest.mark.asyncio
async def test_predict_async_disable_telemetry(classification_model: ClassificationModel, label_names: list[str]):
    """Test async prediction with telemetry disabled"""
    predictions = await classification_model.apredict(["Do you love soup?", "Are cats cute?"], save_telemetry="off")
    assert len(predictions) == 2
    assert predictions[0].prediction_id is None
    assert predictions[1].prediction_id is None
    assert predictions[0].label == 0
    assert predictions[0].label_name == label_names[0]
    assert 0 <= predictions[0].confidence <= 1
    assert predictions[1].label == 1
    assert predictions[1].label_name == label_names[1]
    assert 0 <= predictions[1].confidence <= 1


@pytest.mark.asyncio
async def test_predict_async_with_filters(classification_model: ClassificationModel):
    """Test async prediction with filters"""
    # there are no memories with label 0 and key g2, so we force a wrong prediction
    filtered_prediction = await classification_model.apredict("I love soup", filters=[("key", "==", "g2")])
    assert filtered_prediction.label == 1
    assert filtered_prediction.label_name == "cats"


def test_classification_model_repr(classification_model: ClassificationModel):
    assert repr(classification_model) == dedent(
        """\
        ClassificationModel({
            name: 'test_classification_model',
            num_classes: 2,
            memory_lookup_count: 3,
            weigh_memories: True,
            balance_classes: False,
            memoryset: LabeledMemoryset({
                name: 'test_readonly_memoryset',
                length: 22,
                embedding_model: PretrainedEmbeddingModel({ name: GTE_BASE, embedding_dim: 768, max_seq_length: 8192, num_params: 137M }),
            }),
        })"""
    )


def test_classification_metrics_repr():
    metrics = ClassificationMetrics(
        {
            "coverage": 1.0,
            "accuracy": 0.8333,
            "f1_score": 0.7567,
            "loss": 0.5,
            "anomaly_score_mean": -0.1234,
            "anomaly_score_median": -0.1,
            "anomaly_score_variance": 0.0567,
            "roc_auc": 0.9123,
            "pr_auc": 0.8765,
        }
    )
    assert repr(metrics) == dedent(
        """\
        ClassificationMetrics({
            accuracy: 0.8333,
            f1_score: 0.7567,
            roc_auc: 0.9123,
            pr_auc: 0.8765,
            anomaly_score: -0.1234 ± 0.0567,
        })"""
    )


def test_classification_metrics_repr_without_optional_fields():
    metrics = ClassificationMetrics(
        {
            "coverage": 1.0,
            "accuracy": 0.5,
            "f1_score": 0.5,
            "loss": None,
        }
    )
    assert repr(metrics) == dedent(
        """\
        ClassificationMetrics({
            accuracy: 0.5000,
            f1_score: 0.5000,
        })"""
    )


def test_classification_prediction_repr():
    prediction = ClassificationPrediction(
        None,
        label=0,
        label_name="soup",
        score=None,
        confidence=0.8567,
        anomaly_score=-0.1234,
        memoryset=MagicMock(),
        model=MagicMock(),
        input_value="Do you love soup?",
    )
    assert (
        repr(prediction)
        == "ClassificationPrediction({ label: <soup: 0>, confidence: 0.86, anomaly_score: -0.12, input_value: 'Do you love soup?' })"
    )
