# Copyright (c) 2025 Vertex AI Batch Embeddings API Contributors
#
# Licensed under the MIT License. See LICENSE file in the project root for details.


"""
Tests for the EmbeddingService
"""

from unittest.mock import MagicMock, patch

import pytest

from src.service import EmbeddingService


@pytest.fixture
def mock_storage_client():
    """Fixture for mocking the GCS client."""
    with patch("google.cloud.storage.Client") as mock_client:
        yield mock_client


@pytest.fixture
def mock_aiplatform():
    """Fixture for mocking the AI Platform client."""
    with patch("google.cloud.aiplatform.init") as mock_init:
        with patch("google.cloud.aiplatform.BatchPredictionJob") as mock_job:
            _ = mock_init  # Prevent actual AI Platform initialization
            yield mock_job


@patch.dict(
    "os.environ",
    {
        "GOOGLE_CLOUD_PROJECT": "test-project",
        "GCS_EMBED_INPUT_BUCKET": "test-input-bucket",
        "GCS_EMBED_OUTPUT_BUCKET": "test-output-bucket",
    },
)
def test_trigger_batch_prediction(mock_aiplatform, mock_storage_client):
    """Test that trigger_batch_prediction calls BatchPredictionJob.create with correct params."""
    service = EmbeddingService()

    # Mock the GCS blob object
    mock_blob = MagicMock()
    mock_bucket = MagicMock()
    mock_bucket.blob.return_value = mock_blob
    mock_storage_client.return_value.bucket.return_value = mock_bucket

    # Mock the BatchPredictionJob.create method
    mock_job = MagicMock()
    mock_job.display_name = "test-job"
    mock_job.resource_name = "projects/test-project/locations/us-central1/batchPredictionJobs/123"
    mock_aiplatform.create.return_value = mock_job

    service.trigger_batch_prediction(texts=["text1", "text2"])

    mock_aiplatform.create.assert_called_once()
    _, kwargs = mock_aiplatform.create.call_args
    assert kwargs["instances_format"] == "jsonl"


@patch.dict(
    "os.environ",
    {
        "GOOGLE_CLOUD_PROJECT": "test-project",
        "GCS_EMBED_INPUT_BUCKET": "test-input-bucket",
        "GCS_EMBED_OUTPUT_BUCKET": "test-output-bucket",
    },
)
def test_get_job_status_found(mock_aiplatform, mock_storage_client):
    """Test that get_job_status returns the correct status when a job is found."""
    service = EmbeddingService()

    mock_job = MagicMock()
    mock_job.display_name = "test-job"
    mock_job.resource_name = "projects/test-project/locations/us-central1/batchPredictionJobs/123"
    mock_job.state.name = "JOB_STATE_SUCCEEDED"
    mock_job.create_time.isoformat.return_value = "2023-11-10T12:00:00Z"
    mock_job.start_time.isoformat.return_value = "2023-11-10T12:01:00Z"
    mock_job.end_time.isoformat.return_value = "2023-11-10T12:05:00Z"
    mock_job.output_info.gcs_output_directory = "gs://test-output-bucket/outputs/test-job/"
    mock_job.error = None

    mock_aiplatform.list.return_value = [mock_job]

    status = service.get_job_status(job_name="test-job")

    assert status["job_name"] == "test-job"
    assert status["status"] == "JOB_STATE_SUCCEEDED"


@patch.dict(
    "os.environ",
    {
        "GOOGLE_CLOUD_PROJECT": "test-project",
        "GCS_EMBED_INPUT_BUCKET": "test-input-bucket",
        "GCS_EMBED_OUTPUT_BUCKET": "test-output-bucket",
    },
)
def test_get_job_status_not_found(mock_aiplatform, mock_storage_client):
    """Test that get_job_status raises ValueError when a job is not found."""
    service = EmbeddingService()

    mock_aiplatform.list.return_value = []

    with pytest.raises(ValueError, match="Job 'test-job' not found"):
        service.get_job_status(job_name="test-job")
