api_gemini 0.5.0

//! Batch Mode API for async job-based processing with 50% cost discount.
//!
//! Provides methods for creating batch jobs, polling status, and retrieving results.
//! Batch Mode offers 50% cost discount with 24-hour SLO.
//!
//! **NOTE:** As of 2025-10-11, this implementation uses mock responses. The real
//! Batch API endpoints (e.g., `/v1/batches`) are not yet available in v1beta.
//! Replace mock implementations with real API calls when the endpoints become available.

use crate::
{
  client ::Client,
  error ::Error,
  models ::
  {
    GenerateContentRequest,
    GenerateContentResponse,
    ContentEmbedding,
    Content,
    Part,
    batch ::*,
  },
};
use std::time::{ Duration, SystemTime };

/// API for managing batch jobs with async processing.
#[ derive( Debug ) ]
pub struct BatchApi< 'a >
{
  client : &'a Client,
}

impl< 'a > BatchApi< 'a >
{
  /// Create a new BatchApi instance.
  #[ inline ]
  pub fn new( client : &'a Client ) -> Self
  {
    Self { client }
  }

  /// Create a batch job with inline content generation requests.
  ///
  /// # Arguments
  ///
  /// * `model` - Model name (e.g., "gemini-2.5-flash")
  /// * `requests` - Vec of GenerateContentRequest to process
  ///
  /// # Returns
  ///
  /// Returns a BatchJob with job_id for status polling.
  ///
  /// # Errors
  ///
  /// Returns error if job creation fails or request is invalid.
  pub async fn create_inline(
    &self,
    model : &str,
    requests : Vec< GenerateContentRequest >
  ) -> Result< BatchJob, Error >
  {
    // Use client field to access base_url even in mock implementation
    let _base_url = &self.client.base_url;

    let job_id = format!( "batch_job_{}", uuid::Uuid::new_v4() );
    let request_count = requests.len();

    // Create batch job (mock implementation - replace with real API call)
    let batch_job = BatchJob
    {
      job_id : job_id.clone(),
      state : BatchJobState::Pending,
      model : model.to_string(),
      request_count,
      create_time : Some( SystemTime::now() ),
      expiration_time : Some( SystemTime::now() + Duration::from_secs( 86400 ) ), // 24 hours
      error : None,
    };

    Ok( batch_job )
  }

  /// Get status of a batch job.
  ///
  /// # Arguments
  ///
  /// * `job_id` - The batch job identifier
  ///
  /// # Returns
  ///
  /// Returns BatchJobStatus with current state and progress.
  ///
  /// # Errors
  ///
  /// Returns error if job not found or API call fails.
  pub async fn get_status( &self, job_id : &str ) -> Result< BatchJobStatus, Error >
  {
    // Mock implementation - replace with real API call
    let status = BatchJobStatus
    {
      job_id : job_id.to_string(),
      state : BatchJobState::Running,
      completed_count : 0,
      failed_count : 0,
      update_time : Some( SystemTime::now() ),
      error : None,
    };

    Ok( status )
  }

  /// Wait for batch job completion and retrieve results.
  ///
  /// Polls job status until completion or timeout.
  ///
  /// # Arguments
  ///
  /// * `job_id` - The batch job identifier
  /// * `timeout` - Maximum time to wait
  ///
  /// # Returns
  ///
  /// Returns BatchJobResults with all responses.
  ///
  /// # Errors
  ///
  /// Returns error if timeout reached or job fails.
  pub async fn wait_and_retrieve(
    &self,
    job_id : &str,
    timeout : Duration
  ) -> Result< BatchJobResults, Error >
  {
    let start = SystemTime::now();
    let poll_interval = Duration::from_secs( 5 );

    loop
    {
      let status = self.get_status( job_id ).await?;

      match status.state
      {
        BatchJobState::Succeeded | BatchJobState::PartiallyCompleted =>
        {
          return self.retrieve_results( job_id ).await;
        }
        BatchJobState::Failed =>
        {
          return Err( Error::ApiError(
            status.error.unwrap_or_else( || "Batch job failed".to_string() )
          ) );
        }
        BatchJobState::Cancelled =>
        {
          return Err( Error::ApiError( "Batch job was cancelled".to_string() ) );
        }
        BatchJobState::Pending | BatchJobState::Running =>
        {
          // Check timeout
          if start.elapsed().unwrap_or( Duration::ZERO ) > timeout
          {
            return Err( Error::ApiError( "Batch job timeout".to_string() ) );
          }

          // Wait before next poll
          tokio ::time::sleep( poll_interval ).await;
        }
      }
    }
  }

  /// Retrieve results from a completed batch job.
  ///
  /// # Arguments
  ///
  /// * `job_id` - The batch job identifier
  ///
  /// # Returns
  ///
  /// Returns BatchJobResults with responses and billing info.
  ///
  /// # Errors
  ///
  /// Returns error if results not available or expired.
  async fn retrieve_results( &self, job_id : &str ) -> Result< BatchJobResults, Error >
  {
    // Mock implementation - replace with real API call
    let results = BatchJobResults
    {
      job_id : job_id.to_string(),
      state : BatchJobState::Succeeded,
      responses : vec!
      [
        GenerateContentResponse
        {
          candidates : vec!
          [
            crate ::models::Candidate
            {
              content : Content
              {
                parts : vec!
                [
                  Part
                  {
                    text : Some( "Mock response".to_string() ),
                    ..Default::default()
                  }
                ],
                role : "model".to_string(),
              },
              finish_reason : Some( "STOP".to_string() ),
              safety_ratings : None,
              citation_metadata : None,
              token_count : Some( 10 ),
              index : Some( 0 ),
            }
          ],
          prompt_feedback : None,
          usage_metadata : None,
          grounding_metadata : None,
        }
      ],
      billing_metadata : Some( BatchBillingMetadata
      {
        discount_percentage : 50,
        standard_cost : 0.02,
        discounted_cost : 0.01,
        total_tokens : 100,
      } ),
      retrieve_time : Some( SystemTime::now() ),
    };

    Ok( results )
  }

  /// Cancel a running batch job.
  ///
  /// # Arguments
  ///
  /// * `job_id` - The batch job identifier
  ///
  /// # Errors
  ///
  /// Returns error if job cannot be cancelled or not found.
  pub async fn cancel( &self, job_id : &str ) -> Result< (), Error >
  {
    // Mock implementation - replace with real API call
    let _ = job_id;
    Ok( () )
  }

  /// List all batch jobs.
  ///
  /// # Returns
  ///
  /// Returns BatchJobList with jobs and optional next page token.
  ///
  /// # Errors
  ///
  /// Returns error if list operation fails.
  pub async fn list( &self ) -> Result< BatchJobList, Error >
  {
    self.list_with_page_size( None, None ).await
  }

  /// List batch jobs with pagination token.
  ///
  /// # Arguments
  ///
  /// * `page_token` - Token from previous list response
  ///
  /// # Returns
  ///
  /// Returns next page of batch jobs.
  ///
  /// # Errors
  ///
  /// Returns error if token invalid or list fails.
  pub async fn list_with_token( &self, page_token : &str ) -> Result< BatchJobList, Error >
  {
    self.list_with_page_size( None, Some( page_token.to_string() ) ).await
  }

  /// Internal list implementation with page size and token.
  async fn list_with_page_size(
    &self,
    _page_size : Option< i32 >,
    _page_token : Option< String >
  ) -> Result< BatchJobList, Error >
  {
    // Mock implementation - replace with real API call
    let list = BatchJobList
    {
      jobs : vec![],
      next_page_token : None,
    };

    Ok( list )
  }

  /// Create a batch job for embedding generation.
  ///
  /// # Arguments
  ///
  /// * `model` - Embedding model name (e.g., "text-embedding-004")
  /// * `texts` - Vec of text strings to embed
  ///
  /// # Returns
  ///
  /// Returns BatchJob for polling and result retrieval.
  ///
  /// # Errors
  ///
  /// Returns error if job creation fails.
  pub async fn create_embedding_batch(
    &self,
    model : &str,
    texts : Vec< String >
  ) -> Result< BatchJob, Error >
  {
    let job_id = format!( "batch_embed_{}", uuid::Uuid::new_v4() );
    let request_count = texts.len();

    let batch_job = BatchJob
    {
      job_id : job_id.clone(),
      state : BatchJobState::Pending,
      model : model.to_string(),
      request_count,
      create_time : Some( SystemTime::now() ),
      expiration_time : Some( SystemTime::now() + Duration::from_secs( 86400 ) ),
      error : None,
    };

    Ok( batch_job )
  }

  /// Wait for embedding batch completion and retrieve results.
  ///
  /// # Arguments
  ///
  /// * `job_id` - The batch job identifier
  /// * `timeout` - Maximum time to wait
  ///
  /// # Returns
  ///
  /// Returns BatchEmbeddingResults with all embeddings.
  ///
  /// # Errors
  ///
  /// Returns error if timeout or job fails.
  pub async fn wait_and_retrieve_embeddings(
    &self,
    job_id : &str,
    timeout : Duration
  ) -> Result< BatchEmbeddingResults, Error >
  {
    let start = SystemTime::now();
    let poll_interval = Duration::from_secs( 5 );

    loop
    {
      let status = self.get_status( job_id ).await?;

      match status.state
      {
        BatchJobState::Succeeded | BatchJobState::PartiallyCompleted =>
        {
          return self.retrieve_embedding_results( job_id ).await;
        }
        BatchJobState::Failed =>
        {
          return Err( Error::ApiError(
            status.error.unwrap_or_else( || "Batch job failed".to_string() )
          ) );
        }
        BatchJobState::Cancelled =>
        {
          return Err( Error::ApiError( "Batch job was cancelled".to_string() ) );
        }
        BatchJobState::Pending | BatchJobState::Running =>
        {
          if start.elapsed().unwrap_or( Duration::ZERO ) > timeout
          {
            return Err( Error::ApiError( "Batch job timeout".to_string() ) );
          }

          tokio ::time::sleep( poll_interval ).await;
        }
      }
    }
  }

  /// Retrieve embedding results from completed job.
  async fn retrieve_embedding_results( &self, job_id : &str ) -> Result< BatchEmbeddingResults, Error >
  {
    // Mock implementation - replace with real API call
    let results = BatchEmbeddingResults
    {
      job_id : job_id.to_string(),
      state : BatchJobState::Succeeded,
      embeddings : vec!
      [
        ContentEmbedding
        {
          values : vec![ 0.1, 0.2, 0.3 ],
        }
      ],
      billing_metadata : Some( BatchBillingMetadata
      {
        discount_percentage : 50,
        standard_cost : 0.01,
        discounted_cost : 0.005,
        total_tokens : 50,
      } ),
    };

    Ok( results )
  }
}